title image classification with convmixer sayak paulhttps twitter comrisingsayak date created 20211012 last modified 20211012 description an allconvolutional network applied to patches of images accelerator gpu converted to keras 3 by md awsafur rahmanhttps awsaf49 github io introduction vision transformers vit dosovitskiy et al https arxiv orgabs1612 00593 extract small patches from the input images linearly project them and then apply the transformer vaswani et al https arxiv orgabs1706 03762 blocks the application of vits to image recognition tasks is quickly becoming a promising area of research because vits eliminate the need to have strong inductive biases such as convolutions for modeling locality this presents them as a general computation primititive capable of learning just from the training data with as minimal inductive priors as possible vits yield great downstream performance when trained with proper regularization data augmentation and relatively large datasets in the patches are all you needhttps openreview netpdf idtvhs5y4dnvm paper note at the time of writing it is a submission to the iclr 2022 conference the s extend the idea of using patches to train an allconvolutional network and demonstrate competitive results their architecture namely convmixer uses recipes from the recent isotrophic architectures like vit mlpmixer tolstikhin et al https arxiv orgabs2105 01601 such as using the same depth and resolution across different layers in the network residual connections and so on in this example we will implement the convmixer model and demonstrate its performance on the cifar10 dataset imports hyperparameters to keep run time short we will train the model for only 10 epochs to focus on the core ideas of convmixer we will not use other trainingspecific elements like randaugment cubuk et al https arxiv orgabs1909 13719 if you are interested in learning more about those details please refer to the original paperhttps openreview netpdf idtvhs5y4dnvm load the cifar10 dataset prepare tf data dataset objects our data augmentation pipeline is different from what the s used for the cifar10 dataset which is fine for the purpose of the example note that it s ok to use tf apis for data io and preprocessing with other backends jax torch as it is featurecomplete framework when it comes to data preprocessing convmixer utilities the following figure taken from the original paper depicts the convmixer model https i imgur comyf8actg png convmixer is very similar to the mlpmixer model with the following key differences instead of using fullyconnected layers it uses standard convolution layers instead of layernorm which is typical for vits and mlpmixers it uses batchnorm two types of convolution layers are used in convmixer 1 depthwise convolutions for mixing spatial locations of the images 2 pointwise convolutions which follow the depthwise convolutions for mixing channelwise information across the patches another keypoint is the use of larger kernel sizes to allow a larger receptive field depthwise convolution pointwise convolution convmixer2568 https openreview netpdf idtvhs5y4dnvm the hyperparameter values are taken from the paper extract patch embeddings convmixer blocks classification block the model used in this experiment is termed as convmixer2568 where 256 denotes the number of channels and 8 denotes the depth the resulting model only has 0 8 million parameters model training and evaluation utility code reference https keras ioexamplesvisionimageclassificationwithvisiontransformer train and evaluate model the gap in training and validation performance can be mitigated by using additional regularization techniques nevertheless being able to get to 83 accuracy within 10 epochs with 0 8 million parameters is a strong result visualizing the internals of convmixer we can visualize the patch embeddings and the learned convolution filters recall that each patch embedding and intermediate feature map have the same number of channels 256 in this case this will make our visualization utility easier to implement code reference https bit ly3awirbp first apply minmax normalization to the given weights to avoid isotrophic scaling visualize all the filters we first visualize the learned patch embeddings even though we did not train the network to convergence we can notice that different patches show different patterns some share similarity with others while some are very different these visualizations are more salient with larger image sizes similarly we can visualize the raw convolution kernels this can help us understand the patterns to which a given kernel is receptive first print the indices of the convolution layers that are not pointwise convolutions we see that different filters in the kernel have different locality spans and this pattern is likely to evolve with more training final notes there s been a recent trend on fusing convolutions with other dataagnostic operations like selfattention following works are along this line of research convit d ascoli et al https arxiv orgabs2103 10697 cct hassani et al https arxiv orgabs2104 05704 coatnet dai et al https arxiv orgabs2106 04803 introduction vision transformers vit dosovitskiy et al https arxiv org abs 1612 00593 extract small patches from the input images linearly project them and then apply the transformer vaswani et al https arxiv org abs 1706 03762 blocks the application of vits to image recognition tasks is quickly becoming a promising area of research because vits eliminate the need to have strong inductive biases such as convolutions for modeling locality this presents them as a general computation primititive capable of learning just from the training data with as minimal inductive priors as possible vits yield great downstream performance when trained with proper regularization data augmentation and relatively large datasets in the patches are all you need https openreview net pdf id tvhs5y4dnvm paper note at the time of writing it is a submission to the iclr 2022 conference the s extend the idea of using patches to train an all convolutional network and demonstrate competitive results their architecture namely convmixer uses recipes from the recent isotrophic architectures like vit mlp mixer tolstikhin et al https arxiv org abs 2105 01601 such as using the same depth and resolution across different layers in the network residual connections and so on in this example we will implement the convmixer model and demonstrate its performance on the cifar 10 dataset imports hyperparameters to keep run time short we will train the model for only 10 epochs to focus on the core ideas of convmixer we will not use other training specific elements like randaugment cubuk et al https arxiv org abs 1909 13719 if you are interested in learning more about those details please refer to the original paper https openreview net pdf id tvhs5y4dnvm load the cifar 10 dataset prepare tf data dataset objects our data augmentation pipeline is different from what the s used for the cifar 10 dataset which is fine for the purpose of the example note that it s ok to use tf apis for data i o and preprocessing with other backends jax torch as it is feature complete framework when it comes to data preprocessing convmixer utilities the following figure taken from the original paper depicts the convmixer model https i imgur com yf8actg png convmixer is very similar to the mlp mixer model with the following key differences instead of using fully connected layers it uses standard convolution layers instead of layernorm which is typical for vits and mlp mixers it uses batchnorm two types of convolution layers are used in convmixer 1 depthwise convolutions for mixing spatial locations of the images 2 pointwise convolutions which follow the depthwise convolutions for mixing channel wise information across the patches another keypoint is the use of larger kernel sizes to allow a larger receptive field depthwise convolution residual pointwise convolution convmixer 256 8 https openreview net pdf id tvhs5y4dnvm the hyperparameter values are taken from the paper extract patch embeddings convmixer blocks classification block the model used in this experiment is termed as convmixer 256 8 where 256 denotes the number of channels and 8 denotes the depth the resulting model only has 0 8 million parameters model training and evaluation utility code reference https keras io examples vision image_classification_with_vision_transformer train and evaluate model the gap in training and validation performance can be mitigated by using additional regularization techniques nevertheless being able to get to 83 accuracy within 10 epochs with 0 8 million parameters is a strong result visualizing the internals of convmixer we can visualize the patch embeddings and the learned convolution filters recall that each patch embedding and intermediate feature map have the same number of channels 256 in this case this will make our visualization utility easier to implement code reference https bit ly 3awirbp first apply min max normalization to the given weights to avoid isotrophic scaling visualize all the filters we first visualize the learned patch embeddings even though we did not train the network to convergence we can notice that different patches show different patterns some share similarity with others while some are very different these visualizations are more salient with larger image sizes similarly we can visualize the raw convolution kernels this can help us understand the patterns to which a given kernel is receptive first print the indices of the convolution layers that are not pointwise convolutions taking a kernel from the middle of the network we see that different filters in the kernel have different locality spans and this pattern is likely to evolve with more training final notes there s been a recent trend on fusing convolutions with other data agnostic operations like self attention following works are along this line of research convit d ascoli et al https arxiv org abs 2103 10697 cct hassani et al https arxiv org abs 2104 05704 coatnet dai et al https arxiv org abs 2106 04803
import keras
from keras import layers

import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np



learning_rate = 0.001
weight_decay = 0.0001
batch_size = 128
num_epochs = 10



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

val_indices = int(len(x_train) * val_split)
new_x_train, new_y_train = x_train[val_indices:], y_train[val_indices:]
x_val, y_val = x_train[:val_indices], y_train[:val_indices]

print(f"Training data samples: {len(new_x_train)}")
print(f"Validation data samples: {len(x_val)}")
print(f"Test data samples: {len(x_test)}")



image_size = 32
auto = tf.data.AUTOTUNE

augmentation_layers = [
    keras.layers.RandomCrop(image_size, image_size),
    keras.layers.RandomFlip("horizontal"),
]


def augment_images(images):
    for layer in augmentation_layers:
        images = layer(images, training=True)
    return images


def make_datasets(images, labels, is_train=False):
    dataset = tf.data.Dataset.from_tensor_slices((images, labels))
    if is_train:
        dataset = dataset.shuffle(batch_size * 10)
    dataset = dataset.batch(batch_size)
    if is_train:
        dataset = dataset.map(
            lambda x, y: (augment_images(x), y), num_parallel_calls=auto
        )
    return dataset.prefetch(auto)


train_dataset = make_datasets(new_x_train, new_y_train, is_train=True)
val_dataset = make_datasets(x_val, y_val)
test_dataset = make_datasets(x_test, y_test)




def activation_block(x):
    x = layers.Activation("gelu")(x)
    return layers.BatchNormalization()(x)


def conv_stem(x, filters: int, patch_size: int):
    x = layers.Conv2D(filters, kernel_size=patch_size, strides=patch_size)(x)
    return activation_block(x)


def conv_mixer_block(x, filters: int, kernel_size: int):
    
    x0 = x
    x = layers.DepthwiseConv2D(kernel_size=kernel_size, padding="same")(x)
    x = layers.Add()([activation_block(x), x0])  

    
    x = layers.Conv2D(filters, kernel_size=1)(x)
    x = activation_block(x)

    return x


def get_conv_mixer_256_8(
    image_size=32,
    filters=256,
    depth=8,
    kernel_size=5,
    patch_size=2,
    num_classes=10,
):
    
    inputs = keras.Input((image_size, image_size, 3))
    x = layers.Rescaling(scale=1.0 / 255)(inputs)

    
    x = conv_stem(x, filters, patch_size)

    
    for _ in range(depth):
        x = conv_mixer_block(x, filters, kernel_size)

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

    return keras.Model(inputs, outputs)










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

    model.compile(
        optimizer=optimizer,
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"],
    )

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

    history = model.fit(
        train_dataset,
        validation_data=val_dataset,
        epochs=num_epochs,
        callbacks=[checkpoint_callback],
    )

    model.load_weights(checkpoint_filepath)
    _, accuracy = model.evaluate(test_dataset)
    print(f"Test accuracy: {round(accuracy * 100, 2)}%")

    return history, model




conv_mixer_model = get_conv_mixer_256_8()
history, conv_mixer_model = run_experiment(conv_mixer_model)








def visualization_plot(weights, idx=1):
    
    
    p_min, p_max = weights.min(), weights.max()
    weights = (weights - p_min) / (p_max - p_min)

    
    num_filters = 256
    plt.figure(figsize=(8, 8))

    for i in range(num_filters):
        current_weight = weights[:, :, :, i]
        if current_weight.shape[-1] == 1:
            current_weight = current_weight.squeeze()
        ax = plt.subplot(16, 16, idx)
        ax.set_xticks([])
        ax.set_yticks([])
        plt.imshow(current_weight)
        idx += 1



patch_embeddings = conv_mixer_model.layers[2].get_weights()[0]
visualization_plot(patch_embeddings)





for i, layer in enumerate(conv_mixer_model.layers):
    if isinstance(layer, layers.DepthwiseConv2D):
        if layer.get_config()["kernel_size"] == (5, 5):
            print(i, layer)

idx = 26  

kernel = conv_mixer_model.layers[idx].get_weights()[0]
kernel = np.expand_dims(kernel.squeeze(), axis=2)
visualization_plot(kernel)