title mobilevit a mobilefriendly transformerbased model for image classification sayak paulhttps twitter comrisingsayak date created 20211020 last modified 20231123 description mobilevit for image classification with combined benefits of convolutions and transformers accelerator gpu converted to keras 3 by pavan kumar singhhttps github compksx01 introduction in this example we implement the mobilevit architecture mehta et al https arxiv orgabs2110 02178 which combines the benefits of transformers vaswani et al https arxiv orgabs1706 03762 and convolutions with transformers we can capture longrange dependencies that result in global representations with convolutions we can capture spatial relationships that model locality besides combining the properties of transformers and convolutions the s introduce mobilevit as a generalpurpose mobilefriendly backbone for different image recognition tasks their findings suggest that performancewise mobilevit is better than other models with the same or higher complexity mobilenetv3https arxiv orgabs1905 02244 for example while being efficient on mobile devices imports hyperparameters values are from table 4 mobilevit utilities the mobilevit architecture is comprised of the following blocks strided 3x3 convolutions that process the input image mobilenetv2https arxiv orgabs1801 04381style inverted residual blocks for downsampling the resolution of the intermediate feature maps mobilevit blocks that combine the benefits of transformers and convolutions it is presented in the figure below taken from the original paperhttps arxiv orgabs2110 02178 https i imgur commannhi7 png reference https git iojkgtc reference https keras ioexamplesvisionimageclassificationwithvisiontransformer layer normalization 1 create a multihead attention layer skip connection 1 layer normalization 2 mlp skip connection 2 local projection with convolutions unfold into patches and then pass through transformers fold into convlike featuremaps apply pointwise conv concatenate with the input features fuse the local and global features using a convoluion layer more on the mobilevit block first the feature representations a go through convolution blocks that capture local relationships the expected shape of a single entry here would be h w numchannels then they get unfolded into another vector with shape p n numchannels where p is the area of a small patch and n is h w p so we end up with n nonoverlapping patches this unfolded vector is then passed through a tranformer block that captures global relationships between the patches the output vector b is again folded into a vector of shape h w numchannels resembling a feature map coming out of convolutions vectors a and b are then passed through two more convolutional layers to fuse the local and global representations notice how the spatial resolution of the final vector remains unchanged at this point the s also present an explanation of how the mobilevit block resembles a convolution block of a cnn for more details please refer to the original paper next we combine these blocks together and implement the mobilevit architecture xxs variant the following figure taken from the original paper presents a schematic representation of the architecture https i ibb cosrbvrbnimage png initial convstem mv2 block downsampling with mv2 block first mv2 mobilevit block second mv2 mobilevit block third mv2 mobilevit block classification head dataset preparation we will be using the tfflowershttps www tensorflow orgdatasetscatalogtfflowers dataset to demonstrate the model unlike other transformerbased architectures mobilevit uses a simple augmentation pipeline primarily because it has the properties of a cnn resize to a bigger spatial resolution and take the random crops the s use a multiscale data sampler to help the model learn representations of varied scales in this example we discard this part load and prepare the dataset train a mobilevit xxs model results and tflite conversion with about one million parameters getting to 85 top1 accuracy on 256x256 resolution is a strong result this mobilevit mobile is fully compatible with tensorflow lite tflite and can be converted with the following code serialize the model as a savedmodel convert to tflite this form of quantization is called posttraining dynamicrange quantization in tflite to learn more about different quantization recipes available in tflite and running inference with tflite models check out this official resourcehttps www tensorflow orgliteperformanceposttrainingquantization introduction in this example we implement the mobilevit architecture mehta et al https arxiv org abs 2110 02178 which combines the benefits of transformers vaswani et al https arxiv org abs 1706 03762 and convolutions with transformers we can capture long range dependencies that result in global representations with convolutions we can capture spatial relationships that model locality besides combining the properties of transformers and convolutions the s introduce mobilevit as a general purpose mobile friendly backbone for different image recognition tasks their findings suggest that performance wise mobilevit is better than other models with the same or higher complexity mobilenetv3 https arxiv org abs 1905 02244 for example while being efficient on mobile devices imports hyperparameters values are from table 4 2x2 for the transformer blocks expansion factor for the mobilenetv2 blocks mobilevit utilities the mobilevit architecture is comprised of the following blocks strided 3x3 convolutions that process the input image mobilenetv2 https arxiv org abs 1801 04381 style inverted residual blocks for downsampling the resolution of the intermediate feature maps mobilevit blocks that combine the benefits of transformers and convolutions it is presented in the figure below taken from the original paper https arxiv org abs 2110 02178 https i imgur com mannhi7 png reference https git io jkgtc reference https keras io examples vision image_classification_with_vision_transformer layer normalization 1 create a multi head attention layer skip connection 1 layer normalization 2 mlp skip connection 2 local projection with convolutions unfold into patches and then pass through transformers fold into conv like feature maps apply point wise conv concatenate with the input features fuse the local and global features using a convoluion layer more on the mobilevit block first the feature representations a go through convolution blocks that capture local relationships the expected shape of a single entry here would be h w num_channels then they get unfolded into another vector with shape p n num_channels where p is the area of a small patch and n is h w p so we end up with n non overlapping patches this unfolded vector is then passed through a tranformer block that captures global relationships between the patches the output vector b is again folded into a vector of shape h w num_channels resembling a feature map coming out of convolutions vectors a and b are then passed through two more convolutional layers to fuse the local and global representations notice how the spatial resolution of the final vector remains unchanged at this point the s also present an explanation of how the mobilevit block resembles a convolution block of a cnn for more details please refer to the original paper next we combine these blocks together and implement the mobilevit architecture xxs variant the following figure taken from the original paper presents a schematic representation of the architecture https i ibb co srbvrbn image png initial conv stem mv2 block downsampling with mv2 block first mv2 mobilevit block second mv2 mobilevit block third mv2 mobilevit block classification head dataset preparation we will be using the tf_flowers https www tensorflow org datasets catalog tf_flowers dataset to demonstrate the model unlike other transformer based architectures mobilevit uses a simple augmentation pipeline primarily because it has the properties of a cnn resize to a bigger spatial resolution and take the random crops the s use a multi scale data sampler to help the model learn representations of varied scales in this example we discard this part load and prepare the dataset train a mobilevit xxs model results and tflite conversion with about one million parameters getting to 85 top 1 accuracy on 256x256 resolution is a strong result this mobilevit mobile is fully compatible with tensorflow lite tflite and can be converted with the following code serialize the model as a savedmodel convert to tflite this form of quantization is called post training dynamic range quantization in tflite enable tensorflow lite ops enable tensorflow ops to learn more about different quantization recipes available in tflite and running inference with tflite models check out this official resource https www tensorflow org lite performance post_training_quantization
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf

from keras import layers
import keras as keras

import tensorflow_datasets as tfds

tfds.disable_progress_bar()




patch_size = 4  
image_size = 256
expansion_factor = 2  




def conv_block(x, filters=16, kernel_size=3, strides=2):
    conv_layer = layers.Conv2D(
        filters, kernel_size, strides=strides, activation=keras.activations.swish, padding="same"
    )
    return conv_layer(x)





def inverted_residual_block(x, expanded_channels, output_channels, strides=1):
    m = layers.Conv2D(expanded_channels, 1, padding="same", use_bias=False)(x)
    m = layers.BatchNormalization()(m)
    m = keras.activations.swish(m)

    if strides == 2:
        m = layers.ZeroPadding2D()(m)
    m = layers.DepthwiseConv2D(
        3, strides=strides, padding="same" if strides == 1 else "valid", use_bias=False
    )(m)
    m = layers.BatchNormalization()(m)
    m = keras.activations.swish(m)

    m = layers.Conv2D(output_channels, 1, padding="same", use_bias=False)(m)
    m = layers.BatchNormalization()(m)

    if x.shape[-1] == output_channels and strides == 1:
        return layers.Add()([m, x])
    return m






def mlp(x, hidden_units, dropout_rate):
    for units in hidden_units:
        x = layers.Dense(units, activation=keras.activations.swish)(x)
        x = layers.Dropout(dropout_rate)(x)
    return x


def transformer_block(x, transformer_layers, projection_dim, num_heads=2):
    for _ in range(transformer_layers):
        
        x1 = layers.LayerNormalization(epsilon=1e-6)(x)
        
        attention_output = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=projection_dim, dropout=0.1
        )(x1, x1)
        
        x2 = layers.Add()([attention_output, x])
        
        x3 = layers.LayerNormalization(epsilon=1e-6)(x2)
        
        x3 = mlp(
            x3,
            hidden_units=[x.shape[-1] * 2, x.shape[-1]],
            dropout_rate=0.1,
        )
        
        x = layers.Add()([x3, x2])

    return x


def mobilevit_block(x, num_blocks, projection_dim, strides=1):
    
    local_features = conv_block(x, filters=projection_dim, strides=strides)
    local_features = conv_block(
        local_features, filters=projection_dim, kernel_size=1, strides=strides
    )

    
    num_patches = int((local_features.shape[1] * local_features.shape[2]) / patch_size)
    non_overlapping_patches = layers.Reshape((patch_size, num_patches, projection_dim))(
        local_features
    )
    global_features = transformer_block(
        non_overlapping_patches, num_blocks, projection_dim
    )

    
    folded_feature_map = layers.Reshape((*local_features.shape[1:-1], projection_dim))(
        global_features
    )

    
    folded_feature_map = conv_block(
        folded_feature_map, filters=x.shape[-1], kernel_size=1, strides=strides
    )
    local_global_features = layers.Concatenate(axis=-1)([x, folded_feature_map])

    
    local_global_features = conv_block(
        local_global_features, filters=projection_dim, strides=strides
    )

    return local_global_features







def create_mobilevit(num_classes=5):
    inputs = keras.Input((image_size, image_size, 3))
    x = layers.Rescaling(scale=1.0 / 255)(inputs)

    
    x = conv_block(x, filters=16)
    x = inverted_residual_block(
        x, expanded_channels=16 * expansion_factor, output_channels=16
    )

    
    x = inverted_residual_block(
        x, expanded_channels=16 * expansion_factor, output_channels=24, strides=2
    )
    x = inverted_residual_block(
        x, expanded_channels=24 * expansion_factor, output_channels=24
    )
    x = inverted_residual_block(
        x, expanded_channels=24 * expansion_factor, output_channels=24
    )

    
    x = inverted_residual_block(
        x, expanded_channels=24 * expansion_factor, output_channels=48, strides=2
    )
    x = mobilevit_block(x, num_blocks=2, projection_dim=64)

    
    x = inverted_residual_block(
        x, expanded_channels=64 * expansion_factor, output_channels=64, strides=2
    )
    x = mobilevit_block(x, num_blocks=4, projection_dim=80)

    
    x = inverted_residual_block(
        x, expanded_channels=80 * expansion_factor, output_channels=80, strides=2
    )
    x = mobilevit_block(x, num_blocks=3, projection_dim=96)
    x = conv_block(x, filters=320, kernel_size=1, strides=1)

    
    x = layers.GlobalAvgPool2D()(x)
    outputs = layers.Dense(num_classes, activation="softmax")(x)

    return keras.Model(inputs, outputs)


mobilevit_xxs = create_mobilevit()
mobilevit_xxs.summary()



batch_size = 64
auto = tf.data.AUTOTUNE
resize_bigger = 280
num_classes = 5


def preprocess_dataset(is_training=True):
    def _pp(image, label):
        if is_training:
            
            
            image = tf.image.resize(image, (resize_bigger, resize_bigger))
            image = tf.image.random_crop(image, (image_size, image_size, 3))
            image = tf.image.random_flip_left_right(image)
        else:
            image = tf.image.resize(image, (image_size, image_size))
        label = tf.one_hot(label, depth=num_classes)
        return image, label

    return _pp


def prepare_dataset(dataset, is_training=True):
    if is_training:
        dataset = dataset.shuffle(batch_size * 10)
    dataset = dataset.map(preprocess_dataset(is_training), num_parallel_calls=auto)
    return dataset.batch(batch_size).prefetch(auto)






train_dataset, val_dataset = tfds.load(
    "tf_flowers", split=["train[:90%]", "train[90%:]"], as_supervised=True
)

num_train = train_dataset.cardinality()
num_val = val_dataset.cardinality()
print(f"Number of training examples: {num_train}")
print(f"Number of validation examples: {num_val}")

train_dataset = prepare_dataset(train_dataset, is_training=True)
val_dataset = prepare_dataset(val_dataset, is_training=False)



learning_rate = 0.002
label_smoothing_factor = 0.1
epochs = 30

optimizer = keras.optimizers.Adam(learning_rate=learning_rate)
loss_fn = keras.losses.CategoricalCrossentropy(label_smoothing=label_smoothing_factor)


def run_experiment(epochs=epochs):
    mobilevit_xxs = create_mobilevit(num_classes=num_classes)
    mobilevit_xxs.compile(optimizer=optimizer, loss=loss_fn, metrics=["accuracy"])

    checkpoint_filepath = "/tmp/checkpoint.weights.h5"
    checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        save_best_only=True,
        save_weights_only=True,
    )

    mobilevit_xxs.fit(
        train_dataset,
        validation_data=val_dataset,
        epochs=epochs,
        callbacks=[checkpoint_callback],
    )
    mobilevit_xxs.load_weights(checkpoint_filepath)
    _, accuracy = mobilevit_xxs.evaluate(val_dataset)
    print(f"Validation accuracy: {round(accuracy * 100, 2)}%")
    return mobilevit_xxs


mobilevit_xxs = run_experiment()




tf.saved_model.save(mobilevit_xxs, "mobilevit_xxs")



converter = tf.lite.TFLiteConverter.from_saved_model("mobilevit_xxs")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS,  
    tf.lite.OpsSet.SELECT_TF_OPS,  
]
tflite_model = converter.convert()
open("mobilevit_xxs.tflite", "wb").write(tflite_model)