title compact convolutional transformers sayak paulhttps twitter comrisingsayak date created 20210630 last modified 20230807 description compact convolutional transformers for efficient image classification accelerator gpu converted to keras 3 by muhammad anas razahttps anasrz com guillaume baquiasthttps www linkedin cominguillaumebaquiast478965ba as discussed in the vision transformers vithttps arxiv orgabs2010 11929 paper a transformerbased architecture for vision typically requires a larger dataset than usual as well as a longer pretraining schedule imagenet1khttp imagenet org which has about a million images is considered to fall under the mediumsized data regime with respect to vits this is primarily because unlike cnns vits or a typical transformerbased architecture do not have wellinformed inductive biases such as convolutions for processing images this begs the question can t we combine the benefits of convolution and the benefits of transformers in a single network architecture these benefits include parameterefficiency and selfattention to process longrange and global dependencies interactions between different regions in an image in escaping the big data paradigm with compact transformershttps arxiv orgabs2104 05704 hassani et al present an approach for doing exactly this they proposed the compact convolutional transformer cct architecture in this example we will work on an implementation of cct and we will see how well it performs on the cifar10 dataset if you are unfamiliar with the concept of selfattention or transformers you can read this chapterhttps livebook manning combookdeeplearningwithpythonsecondeditionchapter11r3312 from franois chollet s book deep learning with python this example uses code snippets from another example image classification with vision transformerhttps keras ioexamplesvisionimageclassificationwithvisiontransformer imports hyperparameters and constants load cifar10 dataset the cct tokenizer the first recipe introduced by the cct s is the tokenizer for processing the images in a standard vit images are organized into uniform nonoverlapping patches this eliminates the boundarylevel information present in between different patches this is important for a neural network to effectively exploit the locality information the figure below presents an illustration of how images are organized into patches https i imgur comikbk9oy png we already know that convolutions are quite good at exploiting locality information so based on this the s introduce an allconvolution mininetwork to produce image patches this is our tokenizer after passing the images through our mininetwork the spatial dimensions are flattened to form sequences positional embeddings are optional in cct if we want to use them we can use the layer defined below trim to match the length of the input sequence which might be less than the sequencelength of the layer sequence pooling another recipe introduced in cct is attention pooling or sequence pooling in vit only the feature map corresponding to the class token is pooled and is then used for the subsequent classification task or any other downstream task stochastic depth for regularization stochastic depthhttps arxiv orgabs1603 09382 is a regularization technique that randomly drops a set of layers during inference the layers are kept as they are it is very much similar to dropouthttps jmlr orgpapersv15srivastava14a html but only that it operates on a block of layers rather than individual nodes present inside a layer in cct stochastic depth is used just before the residual blocks of a transformers encoder referred from github com rwightmanpytorchimagemodels mlp for the transformers encoder data augmentation in the original paperhttps arxiv orgabs2104 05704 the s use autoaugmenthttps arxiv orgabs1805 09501 to induce stronger regularization for this example we will be using the standard geometric augmentations like random cropping and flipping note the rescaling layer these layers have predefined inference behavior the final cct model in cct outputs from the transformers encoder are weighted and then passed on to the final taskspecific layer in this example we do classification augment data encode patches apply positional embedding calculate stochastic depth probabilities create multiple layers of the transformer block layer normalization 1 create a multihead attention layer skip connection 1 layer normalization 2 mlp skip connection 2 apply sequence pooling classify outputs create the keras model model training and evaluation let s now visualize the training progress of the model the cct model we just trained has just 0 4 million parameters and it gets us to 79 top1 accuracy within 30 epochs the plot above shows no signs of overfitting as well this means we can train this network for longer perhaps with a bit more regularization and may obtain even better performance this performance can further be improved by additional recipes like cosine decay learning rate schedule other data augmentation techniques like autoaugmenthttps arxiv orgabs1805 09501 mixuphttps arxiv orgabs1710 09412 or cutmixhttps arxiv orgabs1905 04899 with these modifications the s present 95 1 top1 accuracy on the cifar10 dataset the s also present a number of experiments to study how the number of convolution blocks transformers layers etc affect the final performance of ccts for a comparison a vit model takes about 4 7 million parameters and 100 epochs of training to reach a top1 accuracy of 78 22 on the cifar10 dataset you can refer to this notebookhttps colab research google comgistsayakpaul1a80d9f582b044354a1a26c5cb3d69e5imageclassificationwithvisiontransformer ipynb to know about the experimental setup the s also demonstrate the performance of compact convolutional transformers on nlp tasks and they report competitive results there as discussed in the vision transformers vit https arxiv org abs 2010 11929 paper a transformer based architecture for vision typically requires a larger dataset than usual as well as a longer pre training schedule imagenet 1k http imagenet org which has about a million images is considered to fall under the medium sized data regime with respect to vits this is primarily because unlike cnns vits or a typical transformer based architecture do not have well informed inductive biases such as convolutions for processing images this begs the question can t we combine the benefits of convolution and the benefits of transformers in a single network architecture these benefits include parameter efficiency and self attention to process long range and global dependencies interactions between different regions in an image in escaping the big data paradigm with compact transformers https arxiv org abs 2104 05704 hassani et al present an approach for doing exactly this they proposed the compact convolutional transformer cct architecture in this example we will work on an implementation of cct and we will see how well it performs on the cifar 10 dataset if you are unfamiliar with the concept of self attention or transformers you can read this chapter https livebook manning com book deep learning with python second edition chapter 11 r 3 312 from françois chollet s book deep learning with python this example uses code snippets from another example image classification with vision transformer https keras io examples vision image_classification_with_vision_transformer imports hyperparameters and constants load cifar 10 dataset the cct tokenizer the first recipe introduced by the cct s is the tokenizer for processing the images in a standard vit images are organized into uniform non overlapping patches this eliminates the boundary level information present in between different patches this is important for a neural network to effectively exploit the locality information the figure below presents an illustration of how images are organized into patches https i imgur com ikbk9oy png we already know that convolutions are quite good at exploiting locality information so based on this the s introduce an all convolution mini network to produce image patches this is our tokenizer after passing the images through our mini network the spatial dimensions are flattened to form sequences positional embeddings are optional in cct if we want to use them we can use the layer defined below trim to match the length of the input sequence which might be less than the sequence_length of the layer sequence pooling another recipe introduced in cct is attention pooling or sequence pooling in vit only the feature map corresponding to the class token is pooled and is then used for the subsequent classification task or any other downstream task stochastic depth for regularization stochastic depth https arxiv org abs 1603 09382 is a regularization technique that randomly drops a set of layers during inference the layers are kept as they are it is very much similar to dropout https jmlr org papers v15 srivastava14a html but only that it operates on a block of layers rather than individual nodes present inside a layer in cct stochastic depth is used just before the residual blocks of a transformers encoder referred from github com rwightman pytorch image models mlp for the transformers encoder data augmentation in the original paper https arxiv org abs 2104 05704 the s use autoaugment https arxiv org abs 1805 09501 to induce stronger regularization for this example we will be using the standard geometric augmentations like random cropping and flipping note the rescaling layer these layers have pre defined inference behavior the final cct model in cct outputs from the transformers encoder are weighted and then passed on to the final task specific layer in this example we do classification augment data encode patches apply positional embedding calculate stochastic depth probabilities create multiple layers of the transformer block layer normalization 1 create a multi head attention layer skip connection 1 layer normalization 2 mlp skip connection 2 apply sequence pooling classify outputs create the keras model model training and evaluation let s now visualize the training progress of the model the cct model we just trained has just 0 4 million parameters and it gets us to 79 top 1 accuracy within 30 epochs the plot above shows no signs of overfitting as well this means we can train this network for longer perhaps with a bit more regularization and may obtain even better performance this performance can further be improved by additional recipes like cosine decay learning rate schedule other data augmentation techniques like autoaugment https arxiv org abs 1805 09501 mixup https arxiv org abs 1710 09412 or cutmix https arxiv org abs 1905 04899 with these modifications the s present 95 1 top 1 accuracy on the cifar 10 dataset the s also present a number of experiments to study how the number of convolution blocks transformers layers etc affect the final performance of ccts for a comparison a vit model takes about 4 7 million parameters and 100 epochs of training to reach a top 1 accuracy of 78 22 on the cifar 10 dataset you can refer to this notebook https colab research google com gist sayakpaul 1a80d9f582b044354a1a26c5cb3d69e5 image_classification_with_vision_transformer ipynb to know about the experimental setup the s also demonstrate the performance of compact convolutional transformers on nlp tasks and they report competitive results there
from keras import layers
import keras

import matplotlib.pyplot as plt
import numpy as np



positional_emb = True
conv_layers = 2
projection_dim = 128

num_heads = 2
transformer_units = [
    projection_dim,
    projection_dim,
]
transformer_layers = 2
stochastic_depth_rate = 0.1

learning_rate = 0.001
weight_decay = 0.0001
batch_size = 128
num_epochs = 30
image_size = 32



num_classes = 10
input_shape = (32, 32, 3)

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()

y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

print(f"x_train shape: {x_train.shape} - y_train shape: {y_train.shape}")
print(f"x_test shape: {x_test.shape} - y_test shape: {y_test.shape}")




class CCTTokenizer(layers.Layer):
    def __init__(
        self,
        kernel_size=3,
        stride=1,
        padding=1,
        pooling_kernel_size=3,
        pooling_stride=2,
        num_conv_layers=conv_layers,
        num_output_channels=[64, 128],
        positional_emb=positional_emb,
        **kwargs,
    ):
        super().__init__(**kwargs)

        
        self.conv_model = keras.Sequential()
        for i in range(num_conv_layers):
            self.conv_model.add(
                layers.Conv2D(
                    num_output_channels[i],
                    kernel_size,
                    stride,
                    padding="valid",
                    use_bias=False,
                    activation="relu",
                    kernel_initializer="he_normal",
                )
            )
            self.conv_model.add(layers.ZeroPadding2D(padding))
            self.conv_model.add(
                layers.MaxPooling2D(pooling_kernel_size, pooling_stride, "same")
            )

        self.positional_emb = positional_emb

    def call(self, images):
        outputs = self.conv_model(images)
        
        
        reshaped = keras.ops.reshape(
            outputs,
            (
                -1,
                keras.ops.shape(outputs)[1] * keras.ops.shape(outputs)[2],
                keras.ops.shape(outputs)[-1],
            ),
        )
        return reshaped





class PositionEmbedding(keras.layers.Layer):
    def __init__(
        self,
        sequence_length,
        initializer="glorot_uniform",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if sequence_length is None:
            raise ValueError(
                "`sequence_length` must be an Integer, received `None`."
            )
        self.sequence_length = int(sequence_length)
        self.initializer = keras.initializers.get(initializer)

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "sequence_length": self.sequence_length,
                "initializer": keras.initializers.serialize(self.initializer),
            }
        )
        return config

    def build(self, input_shape):
        feature_size = input_shape[-1]
        self.position_embeddings = self.add_weight(
            name="embeddings",
            shape=[self.sequence_length, feature_size],
            initializer=self.initializer,
            trainable=True,
        )

        super().build(input_shape)

    def call(self, inputs, start_index=0):
        shape = keras.ops.shape(inputs)
        feature_length = shape[-1]
        sequence_length = shape[-2]
        
        
        position_embeddings = keras.ops.convert_to_tensor(
            self.position_embeddings
        )
        position_embeddings = keras.ops.slice(
            position_embeddings,
            (start_index, 0),
            (sequence_length, feature_length),
        )
        return keras.ops.broadcast_to(position_embeddings, shape)

    def compute_output_shape(self, input_shape):
        return input_shape





class SequencePooling(layers.Layer):
    def __init__(self):
        super().__init__()
        self.attention = layers.Dense(1)

    def call(self, x):
        attention_weights = keras.ops.softmax(self.attention(x), axis=1)
        attention_weights = keras.ops.transpose(
            attention_weights, axes=(0, 2, 1)
        )
        weighted_representation = keras.ops.matmul(attention_weights, x)
        return keras.ops.squeeze(weighted_representation, -2)






class StochasticDepth(layers.Layer):
    def __init__(self, drop_prop, **kwargs):
        super().__init__(**kwargs)
        self.drop_prob = drop_prop
        self.seed_generator = keras.random.SeedGenerator(1337)

    def call(self, x, training=None):
        if training:
            keep_prob = 1 - self.drop_prob
            shape = (keras.ops.shape(x)[0],) + (1,) * (len(x.shape) - 1)
            random_tensor = keep_prob + keras.random.uniform(
                shape, 0, 1, seed=self.seed_generator
            )
            random_tensor = keras.ops.floor(random_tensor)
            return (x / keep_prob) * random_tensor
        return x





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





data_augmentation = keras.Sequential(
    [
        layers.Rescaling(scale=1.0 / 255),
        layers.RandomCrop(image_size, image_size),
        layers.RandomFlip("horizontal"),
    ],
    name="data_augmentation",
)




def create_cct_model(
    image_size=image_size,
    input_shape=input_shape,
    num_heads=num_heads,
    projection_dim=projection_dim,
    transformer_units=transformer_units,
):
    inputs = layers.Input(input_shape)

    
    augmented = data_augmentation(inputs)

    
    cct_tokenizer = CCTTokenizer()
    encoded_patches = cct_tokenizer(augmented)

    
    if positional_emb:
        sequence_length = encoded_patches.shape[1]
        encoded_patches += PositionEmbedding(sequence_length=sequence_length)(
            encoded_patches
        )

    
    dpr = [x for x in np.linspace(0, stochastic_depth_rate, transformer_layers)]

    
    for i in range(transformer_layers):
        
        x1 = layers.LayerNormalization(epsilon=1e-5)(encoded_patches)

        
        attention_output = layers.MultiHeadAttention(
            num_heads=num_heads, key_dim=projection_dim, dropout=0.1
        )(x1, x1)

        
        attention_output = StochasticDepth(dpr[i])(attention_output)
        x2 = layers.Add()([attention_output, encoded_patches])

        
        x3 = layers.LayerNormalization(epsilon=1e-5)(x2)

        
        x3 = mlp(x3, hidden_units=transformer_units, dropout_rate=0.1)

        
        x3 = StochasticDepth(dpr[i])(x3)
        encoded_patches = layers.Add()([x3, x2])

    
    representation = layers.LayerNormalization(epsilon=1e-5)(encoded_patches)
    weighted_representation = SequencePooling()(representation)

    
    logits = layers.Dense(num_classes)(weighted_representation)
    
    model = keras.Model(inputs=inputs, outputs=logits)
    return model





def run_experiment(model):
    optimizer = keras.optimizers.AdamW(learning_rate=0.001, weight_decay=0.0001)

    model.compile(
        optimizer=optimizer,
        loss=keras.losses.CategoricalCrossentropy(
            from_logits=True, label_smoothing=0.1
        ),
        metrics=[
            keras.metrics.CategoricalAccuracy(name="accuracy"),
            keras.metrics.TopKCategoricalAccuracy(5, name="top-5-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,
    )

    history = model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_split=0.1,
        callbacks=[checkpoint_callback],
    )

    model.load_weights(checkpoint_filepath)
    _, accuracy, top_5_accuracy = model.evaluate(x_test, y_test)
    print(f"Test accuracy: {round(accuracy * 100, 2)}%")
    print(f"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%")

    return history


cct_model = create_cct_model()
history = run_experiment(cct_model)



plt.plot(history.history["loss"], label="train_loss")
plt.plot(history.history["val_loss"], label="val_loss")
plt.xlabel("Epochs")
plt.ylabel("Loss")
plt.title("Train and Validation Losses Over Epochs", fontsize=14)
plt.legend()
plt.grid()
plt.show()