title image classification with modern mlp models khalid salamahttps www linkedin cominkhalidsalama24403144 converted to keras 3 by guillaume baquiasthttps www linkedin cominguillaumebaquiast478965ba divyasreepathttps github comdivyashreepathihalli date created 20210530 last modified 20230803 description implementing the mlpmixer fnet and gmlp models for cifar100 image classification accelerator gpu introduction this example implements three modern attentionfree multilayer perceptron mlp based models for image classification demonstrated on the cifar100 dataset 1 the mlpmixerhttps arxiv orgabs2105 01601 model by ilya tolstikhin et al based on two types of mlps 3 the fnethttps arxiv orgabs2105 03824 model by james leethorp et al based on unparameterized fourier transform 2 the gmlphttps arxiv orgabs2105 08050 model by hanxiao liu et al based on mlp with gating the purpose of the example is not to compare between these models as they might perform differently on different datasets with welltuned hyperparameters rather it is to show simple implementations of their main building blocks setup prepare the data configure the hyperparameters build a classification model we implement a method that builds a classifier given the processing blocks augment data create patches encode patches to generate a batchsize numpatches embeddingdim tensor process x using the module blocks apply global average pooling to generate a batchsize embeddingdim representation tensor apply dropout compute logits outputs create the keras model define an experiment we implement a utility function to compile train and evaluate a given model create adam optimizer with weight decay compile the model create a learning rate scheduler callback create an early stopping callback fit the model return history to plot learning curves use data augmentation compute the mean and the variance of the training data for normalization implement patch extraction as a layer implement position embedding as a layer trim to match the length of the input sequence which might be less than the sequencelength of the layer the mlpmixer model the mlpmixer is an architecture based exclusively on multilayer perceptrons mlps that contains two types of mlp layers 1 one applied independently to image patches which mixes the perlocation features 2 the other applied across patches along channels which mixes spatial information this is similar to a depthwise separable convolution based modelhttps arxiv orgpdf1610 02357 pdf such as the xception model but with two chained dense transforms no max pooling and layer normalization instead of batch normalization implement the mlpmixer module apply layer normalization transpose inputs from numbatches numpatches hiddenunits to numbatches hiddenunits numpatches apply mlp1 on each channel independently transpose mlp1outputs from numbatches hiddendim numpatches to numbatches numpatches hiddenunits add skip connection apply layer normalization apply mlp2 on each patch independtenly add skip connection build train and evaluate the mlpmixer model note that training the model with the current settings on a v100 gpus takes around 8 seconds per epoch the mlpmixer model tends to have much less number of parameters compared to convolutional and transformerbased models which leads to less training and serving computational cost as mentioned in the mlpmixerhttps arxiv orgabs2105 01601 paper when pretrained on large datasets or with modern regularization schemes the mlpmixer attains competitive scores to stateoftheart models you can obtain better results by increasing the embedding dimensions increasing the number of mixer blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes the fnet model the fnet uses a similar block to the transformer block however fnet replaces the selfattention layer in the transformer block with a parameterfree 2d fourier transformation layer 1 one 1d fourier transform is applied along the patches 2 one 1d fourier transform is applied along the channels implement the fnet module apply fourier transformations add skip connection apply layer normalization apply feedfowrad network add skip connection apply layer normalization build train and evaluate the fnet model note that training the model with the current settings on a v100 gpus takes around 8 seconds per epoch as shown in the fnethttps arxiv orgabs2105 03824 paper better results can be achieved by increasing the embedding dimensions increasing the number of fnet blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes the fnet scales very efficiently to long inputs runs much faster than attentionbased transformer models and produces competitive accuracy results the gmlp model the gmlp is a mlp architecture that features a spatial gating unit sgu the sgu enables crosspatch interactions across the spatial channel dimension by 1 transforming the input spatially by applying linear projection across patches along channels 2 applying elementwise multiplication of the input and its spatial transformation implement the gmlp module split x along the channel dimensions tensors u and v will in the shape of batchsize numpatchs embeddingdim apply layer normalization apply spatial projection apply elementwise multiplication apply layer normalization apply the first channel projection xprojected shape batchsize numpatches embeddingdim 2 apply the spatial gating unit xspatial shape batchsize numpatches embeddingdim apply the second channel projection xprojected shape batchsize numpatches embeddingdim add skip connection build train and evaluate the gmlp model note that training the model with the current settings on a v100 gpus takes around 9 seconds per epoch as shown in the gmlphttps arxiv orgabs2105 08050 paper better results can be achieved by increasing the embedding dimensions increasing the number of gmlp blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes note that the paper used advanced regularization strategies such as mixup and cutmix as well as autoaugment introduction this example implements three modern attention free multi layer perceptron mlp based models for image classification demonstrated on the cifar 100 dataset 1 the mlp mixer https arxiv org abs 2105 01601 model by ilya tolstikhin et al based on two types of mlps 3 the fnet https arxiv org abs 2105 03824 model by james lee thorp et al based on unparameterized fourier transform 2 the gmlp https arxiv org abs 2105 08050 model by hanxiao liu et al based on mlp with gating the purpose of the example is not to compare between these models as they might perform differently on different datasets with well tuned hyperparameters rather it is to show simple implementations of their main building blocks setup prepare the data configure the hyperparameters we ll resize input images to this size size of the patches to be extracted from the input images size of the data array number of hidden units number of blocks build a classification model we implement a method that builds a classifier given the processing blocks augment data create patches encode patches to generate a batch_size num_patches embedding_dim tensor process x using the module blocks apply global average pooling to generate a batch_size embedding_dim representation tensor apply dropout compute logits outputs create the keras model define an experiment we implement a utility function to compile train and evaluate a given model create adam optimizer with weight decay compile the model create a learning rate scheduler callback create an early stopping callback fit the model return history to plot learning curves use data augmentation compute the mean and the variance of the training data for normalization implement patch extraction as a layer implement position embedding as a layer trim to match the length of the input sequence which might be less than the sequence_length of the layer the mlp mixer model the mlp mixer is an architecture based exclusively on multi layer perceptrons mlps that contains two types of mlp layers 1 one applied independently to image patches which mixes the per location features 2 the other applied across patches along channels which mixes spatial information this is similar to a depthwise separable convolution based model https arxiv org pdf 1610 02357 pdf such as the xception model but with two chained dense transforms no max pooling and layer normalization instead of batch normalization implement the mlp mixer module apply layer normalization transpose inputs from num_batches num_patches hidden_units to num_batches hidden_units num_patches apply mlp1 on each channel independently transpose mlp1_outputs from num_batches hidden_dim num_patches to num_batches num_patches hidden_units add skip connection apply layer normalization apply mlp2 on each patch independtenly add skip connection build train and evaluate the mlp mixer model note that training the model with the current settings on a v100 gpus takes around 8 seconds per epoch the mlp mixer model tends to have much less number of parameters compared to convolutional and transformer based models which leads to less training and serving computational cost as mentioned in the mlp mixer https arxiv org abs 2105 01601 paper when pre trained on large datasets or with modern regularization schemes the mlp mixer attains competitive scores to state of the art models you can obtain better results by increasing the embedding dimensions increasing the number of mixer blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes the fnet model the fnet uses a similar block to the transformer block however fnet replaces the self attention layer in the transformer block with a parameter free 2d fourier transformation layer 1 one 1d fourier transform is applied along the patches 2 one 1d fourier transform is applied along the channels implement the fnet module apply fourier transformations add skip connection apply layer normalization apply feedfowrad network add skip connection apply layer normalization build train and evaluate the fnet model note that training the model with the current settings on a v100 gpus takes around 8 seconds per epoch as shown in the fnet https arxiv org abs 2105 03824 paper better results can be achieved by increasing the embedding dimensions increasing the number of fnet blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes the fnet scales very efficiently to long inputs runs much faster than attention based transformer models and produces competitive accuracy results the gmlp model the gmlp is a mlp architecture that features a spatial gating unit sgu the sgu enables cross patch interactions across the spatial channel dimension by 1 transforming the input spatially by applying linear projection across patches along channels 2 applying element wise multiplication of the input and its spatial transformation implement the gmlp module split x along the channel dimensions tensors u and v will in the shape of batch_size num_patchs embedding_dim apply layer normalization apply spatial projection apply element wise multiplication apply layer normalization apply the first channel projection x_projected shape batch_size num_patches embedding_dim 2 apply the spatial gating unit x_spatial shape batch_size num_patches embedding_dim apply the second channel projection x_projected shape batch_size num_patches embedding_dim add skip connection build train and evaluate the gmlp model note that training the model with the current settings on a v100 gpus takes around 9 seconds per epoch as shown in the gmlp https arxiv org abs 2105 08050 paper better results can be achieved by increasing the embedding dimensions increasing the number of gmlp blocks and training the model for longer you may also try to increase the size of the input images and use different patch sizes note that the paper used advanced regularization strategies such as mixup and cutmix as well as autoaugment
import numpy as np
import keras
from keras import layers



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

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

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}")



weight_decay = 0.0001
batch_size = 128
num_epochs = 50
dropout_rate = 0.2
image_size = 64  
patch_size = 8  
num_patches = (image_size // patch_size) ** 2  
embedding_dim = 256  
num_blocks = 4  

print(f"Image size: {image_size} X {image_size} = {image_size ** 2}")
print(f"Patch size: {patch_size} X {patch_size} = {patch_size ** 2} ")
print(f"Patches per image: {num_patches}")
print(f"Elements per patch (3 channels): {(patch_size ** 2) * 3}")




def build_classifier(blocks, positional_encoding=False):
    inputs = layers.Input(shape=input_shape)
    
    augmented = data_augmentation(inputs)
    
    patches = Patches(patch_size)(augmented)
    
    x = layers.Dense(units=embedding_dim)(patches)
    if positional_encoding:
        x = x + PositionEmbedding(sequence_length=num_patches)(x)
    
    x = blocks(x)
    
    representation = layers.GlobalAveragePooling1D()(x)
    
    representation = layers.Dropout(rate=dropout_rate)(representation)
    
    logits = layers.Dense(num_classes)(representation)
    
    return keras.Model(inputs=inputs, outputs=logits)





def run_experiment(model):
    
    optimizer = keras.optimizers.AdamW(
        learning_rate=learning_rate,
        weight_decay=weight_decay,
    )
    
    model.compile(
        optimizer=optimizer,
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[
            keras.metrics.SparseCategoricalAccuracy(name="acc"),
            keras.metrics.SparseTopKCategoricalAccuracy(5, name="top5-acc"),
        ],
    )
    
    reduce_lr = keras.callbacks.ReduceLROnPlateau(
        monitor="val_loss", factor=0.5, patience=5
    )
    
    early_stopping = keras.callbacks.EarlyStopping(
        monitor="val_loss", patience=10, restore_best_weights=True
    )
    
    history = model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_split=0.1,
        callbacks=[early_stopping, reduce_lr],
    )

    _, 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




data_augmentation = keras.Sequential(
    [
        layers.Normalization(),
        layers.Resizing(image_size, image_size),
        layers.RandomFlip("horizontal"),
        layers.RandomZoom(height_factor=0.2, width_factor=0.2),
    ],
    name="data_augmentation",
)

data_augmentation.layers[0].adapt(x_train)





class Patches(layers.Layer):
    def __init__(self, patch_size, **kwargs):
        super().__init__(**kwargs)
        self.patch_size = patch_size

    def call(self, x):
        patches = keras.ops.image.extract_patches(x, self.patch_size)
        batch_size = keras.ops.shape(patches)[0]
        num_patches = keras.ops.shape(patches)[1] * keras.ops.shape(patches)[2]
        patch_dim = keras.ops.shape(patches)[3]
        out = keras.ops.reshape(patches, (batch_size, num_patches, patch_dim))
        return out





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 MLPMixerLayer(layers.Layer):
    def __init__(
        self, num_patches, hidden_units, dropout_rate, *args, **kwargs
    ):
        super().__init__(*args, **kwargs)

        self.mlp1 = keras.Sequential(
            [
                layers.Dense(units=num_patches, activation="gelu"),
                layers.Dense(units=num_patches),
                layers.Dropout(rate=dropout_rate),
            ]
        )
        self.mlp2 = keras.Sequential(
            [
                layers.Dense(units=num_patches, activation="gelu"),
                layers.Dense(units=hidden_units),
                layers.Dropout(rate=dropout_rate),
            ]
        )
        self.normalize = layers.LayerNormalization(epsilon=1e-6)

    def build(self, input_shape):
        return super().build(input_shape)

    def call(self, inputs):
        
        x = self.normalize(inputs)
        
        x_channels = keras.ops.transpose(x, axes=(0, 2, 1))
        
        mlp1_outputs = self.mlp1(x_channels)
        
        mlp1_outputs = keras.ops.transpose(mlp1_outputs, axes=(0, 2, 1))
        
        x = mlp1_outputs + inputs
        
        x_patches = self.normalize(x)
        
        mlp2_outputs = self.mlp2(x_patches)
        
        x = x + mlp2_outputs
        return x




mlpmixer_blocks = keras.Sequential(
    [
        MLPMixerLayer(num_patches, embedding_dim, dropout_rate)
        for _ in range(num_blocks)
    ]
)
learning_rate = 0.005
mlpmixer_classifier = build_classifier(mlpmixer_blocks)
history = run_experiment(mlpmixer_classifier)








class FNetLayer(layers.Layer):
    def __init__(self, embedding_dim, dropout_rate, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.ffn = keras.Sequential(
            [
                layers.Dense(units=embedding_dim, activation="gelu"),
                layers.Dropout(rate=dropout_rate),
                layers.Dense(units=embedding_dim),
            ]
        )

        self.normalize1 = layers.LayerNormalization(epsilon=1e-6)
        self.normalize2 = layers.LayerNormalization(epsilon=1e-6)

    def call(self, inputs):
        
        real_part = inputs
        im_part = keras.ops.zeros_like(inputs)
        x = keras.ops.fft2((real_part, im_part))[0]
        
        x = x + inputs
        
        x = self.normalize1(x)
        
        x_ffn = self.ffn(x)
        
        x = x + x_ffn
        
        return self.normalize2(x)




fnet_blocks = keras.Sequential(
    [FNetLayer(embedding_dim, dropout_rate) for _ in range(num_blocks)]
)
learning_rate = 0.001
fnet_classifier = build_classifier(fnet_blocks, positional_encoding=True)
history = run_experiment(fnet_classifier)








class gMLPLayer(layers.Layer):
    def __init__(
        self, num_patches, embedding_dim, dropout_rate, *args, **kwargs
    ):
        super().__init__(*args, **kwargs)

        self.channel_projection1 = keras.Sequential(
            [
                layers.Dense(units=embedding_dim * 2, activation="gelu"),
                layers.Dropout(rate=dropout_rate),
            ]
        )

        self.channel_projection2 = layers.Dense(units=embedding_dim)

        self.spatial_projection = layers.Dense(
            units=num_patches, bias_initializer="Ones"
        )

        self.normalize1 = layers.LayerNormalization(epsilon=1e-6)
        self.normalize2 = layers.LayerNormalization(epsilon=1e-6)

    def spatial_gating_unit(self, x):
        
        
        u, v = keras.ops.split(x, indices_or_sections=2, axis=2)
        
        v = self.normalize2(v)
        
        v_channels = keras.ops.transpose(v, axes=(0, 2, 1))
        v_projected = self.spatial_projection(v_channels)
        v_projected = keras.ops.transpose(v_projected, axes=(0, 2, 1))
        
        return u * v_projected

    def call(self, inputs):
        
        x = self.normalize1(inputs)
        
        x_projected = self.channel_projection1(x)
        
        x_spatial = self.spatial_gating_unit(x_projected)
        
        x_projected = self.channel_projection2(x_spatial)
        
        return x + x_projected




gmlp_blocks = keras.Sequential(
    [
        gMLPLayer(num_patches, embedding_dim, dropout_rate)
        for _ in range(num_blocks)
    ]
)
learning_rate = 0.003
gmlp_classifier = build_classifier(gmlp_blocks)
history = run_experiment(gmlp_classifier)