Package: vit-tf2
vit_tf2This is a package that implements the ViT model based on Tensorflow. The ViT was proposed in the paper "An image is worth 16x16 words: transformers for image recognition at scale". This package uses pre trained weights on the imagenet21K and imagenet2012 datasets, which are in. npz format.◈ PreconditionsPython >= 3.7Tensorflow >= 2.9Q1: What can you do with this package？Build a pre trained standard specification ViT model.Customize and build any specification ViT model to suit your task.Q2: How to build a pre trained ViT？Quickly build a pre trained ViTB16from vit_tf2.vit import ViT_B16
vit = ViT_B16()The pre trained ViT has 4 configurations: ViT_B16, ViT_B32, ViT_L16 and ViT_L32.configpatch sizehiddem dimmlp dimattention headsencoder depthViT_B1616×1676830721212ViT_B3232×3276830721212ViT_L1616×16102440961624ViT_L3232×32102440961624The "imagenet21k" and "imagenet21k+imagenet2012" are slightly different, as shown in the table below.datasetimage sizeclassespre logitsknown labelsimagenet21k22421843TrueFalseimagenet21k+imagenet20123841000FalseTrueBuild ViTB16 with differernt pre trained weights.from vit_tf2.vit import ViT_B16
vit_1 = ViT_B16(weights = "imagenet21k")
vit_2 = ViT_B16(weights="imagenet21k+imagenet2012")Build ViTB16 without pre trained weightsfrom vit_tf2.vit import ViT_B16
vit = ViT_B16(pre_trained=False)The pre training weights file will be downloaded to C:\Users\user_name\. Keras\weights when "pre_trained = True".Build pre trained ViTB32 with custom parametersfrom vit_tf2.vit import ViT_B32
vit = ViT_B32(
    image_size = 128,
    num_classes = 12, 
    pre_logits = False,
    weights = "imagenet21k",
    )When you change some model parameters and some layers change, these layers will not load pre trained weights, the unchanged layers will still load pre trained weights. You can use loading_summary() to view specific information.vit.loading_summary()
>>
Model: "ViT-B-32-128"
-----------------------------------------------------------------
layers                             load weights inf
=================================================================
patch_embedding                    loaded

add_cls_token                      loaded - imagenet

position_embedding                 not loaded - mismatch

transformer_block_0                loaded - imagenet

transformer_block_1                loaded - imagenet

transformer_block_2                loaded - imagenet

transformer_block_3                loaded - imagenet

transformer_block_4                loaded - imagenet

transformer_block_5                loaded - imagenet

transformer_block_6                loaded - imagenet

transformer_block_7                loaded - imagenet

transformer_block_8                loaded - imagenet

transformer_block_9                loaded - imagenet

transformer_block_10               loaded - imagenet

transformer_block_11               loaded - imagenet

layer_norm                         loaded - imagenet

mlp_head                           not loaded - mismatch
=================================================================Q3: How to build a custom ViT？Instantiating ViT classes to build custom ViT modelsfrom vit_tf2.vit import ViT
vit = ViT(
    image_size = 128,
    patch_size = 36,
    num_classes = 1,
    hidden_dim = 128,
    mlp_dim = 512,
    atten_heads = 32,
    encoder_depth = 4,
    dropout_rate = 0.1,
    activation = "sigmoid",
    pre_logits = True,
    include_mlp_head = True,
    )
vit.summary()

>>
Model: "ViT-CUSTOM_SIZE-36-128"
_________________________________________________________________
 Layer (type)                Output Shape              Param #
=================================================================
 patch_embedding (PatchEmbed  (None, 9, 128)           497792
 ding)

 add_cls_token (AddCLSToken)  (None, 10, 128)          128

 position_embedding (AddPosi  (None, 10, 128)          1280
 tionEmbedding)

 transformer_block_0 (Transf  (None, 10, 128)          198272
 ormerEncoder)

 transformer_block_1 (Transf  (None, 10, 128)          198272
 ormerEncoder)

 transformer_block_2 (Transf  (None, 10, 128)          198272
 ormerEncoder)

 transformer_block_3 (Transf  (None, 10, 128)          198272
 ormerEncoder)

 layer_norm (LayerNormalizat  (None, 10, 128)          256
 ion)

 extract_token (Lambda)      (None, 128)               0

 pre_logits (Dense)          (None, 128)               16512

 mlp_head (Dense)            (None, 1)                 129

=================================================================
Total params: 1,309,185
Trainable params: 1,309,185
Non-trainable params: 0
_________________________________________________________________==========================It should be noted that "hidden_dim" should be divisible by "atten_heads". It is best to set "image_size" size that can be evenly divided by "patch_size".Load pre trained weights for custom modelfrom vit_tf2 import utils, vit
vit_custom = vit.ViT(
    image_size=128,
    patch_size=8,
    encoder_depth=4
    )
utils.load_imgnet_weights(vit_custom, "ViT-B_16_imagenet21k.npz")
vit_custom.loading_summary()

>>
Model: "ViT-CUSTOM_SIZE-8-128"
-----------------------------------------------------------------
layers                             load weights inf
=================================================================
patch_embedding                    mismatch

add_cls_token                      loaded - imagenet

position_embedding                 not loaded - mismatch

transformer_block_0                loaded - imagenet

transformer_block_1                loaded - imagenet

transformer_block_2                loaded - imagenet

transformer_block_3                loaded - imagenet

layer_norm                         loaded - imagenet

pre_logits                         loaded - imagenet

mlp_head                           not loaded - mismatch
=================================================================Q4: Fine tuning or image classification on pre trained ViT ？Fine tuning pre trained ViTfrom vit_tf2.vit import ViT_L16

# Set parameters
IMAGE_SIZE = ...
NUM_CLASSES = ...
ACTIVATION = ...
...

# build ViT
vit = ViT_B32(
    image_size = IMAGE_SIZE,
    num_classes = NUM_CLASSES, 
    activation = ACTIVATION,
    )

# Compiling ViT
vit.compile(
    optimizer = ...,
    loss = ...,
    metrics = ...
    )

# Define train, valid and test data
train_generator = ...
valid_generator = ...
test_generator  = ...

# fine tuning ViT
vit.fit(
    x = train_generator ,
    validation_data = valid_generator ,
    steps_per_epoch = ...,
    validation_steps = ...,
    )

# testing
vit.evaluate(x = test_generator, steps=...)Applying  pre trained ViT for Image Classificationfrom vit_tf2 import vit
from vit_tf2 import utils

# Get pre-trained vitb16
vit_model = vit.ViT_B16(weights="imagenet21k+imagenet2012")

# Load a picture
img = utils.read_img("test.jpg", resize=vit_model.image_size)
img = img.reshape((1,*vit_model.image_size,3))

# Classifying
y = vit_model.predict(img)
classes = utils.get_imagenet2012_classes()
print(classes[y[0].argmax()])It should be noted that as there is currently no label for "imagenet21k", please use "imagenet21k+imagenet2012" when applying pre trained ViT. Both "imagenet21k" and "imagenet21k+imagenet2012" are available during the fine-tuning stage.