<jupyter_start><jupyter_text>Example illustrating how to use the MLP-Mixer APIThis example also shows the performance of the MLP-Mixer architecture on the CIFAR-10 dataset<jupyter_code>import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from mlxtend.plotting import plot_confusion_matrix
from sklearn.metrics import confusion_matrix
from tensorflow.keras.layers import *
from tensorflow.keras.layers.experimental.preprocessing import (
    RandomFlip,
    RandomRotation,
    RandomZoom,
)
from tensorflow.keras.models import Model
from mlp_utils import *

# Uncomment line below to run on CPU
# tf.config.set_visible_devices([], "GPU")
# Load the dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()

# normalize pixel values
x_train = x_train / 255
x_test = x_test / 255

# create validation data from test data
val_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))

# Define the classes
classes = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
]
def create_model():
    dims = 512
    k, s = (4, 4)  # (kernel, strides)
    depth = 4

    inputs = Input((32, 32, 3))

    x = tf.keras.Sequential(
        [
            RandomFlip(),
            RandomRotation(factor=0.03),
            RandomZoom(height_factor=0.2, width_factor=0.2),
        ],
        name="data_augmentation",
    )(inputs)

    x = CreatePatches(k, s)(x)
    x = PerPatchFullyConnected(dims)(x)

    for _ in range(depth):
        x = MLPBlock()(x)

    # x = Projection()(x)
    x = GlobalAveragePooling1D()(x)

    # classification layer
    x = GaussianDropout(0.9)(x)
    output = Dense(len(classes), activation="softmax", kernel_regularizer="l2")(x)

    return Model(inputs=inputs, outputs=output)


model = create_model()
print(model.summary())<jupyter_output>Model: "model"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 32, 32, 3)]       0         
_________________________________________________________________
data_augmentation (Sequentia (None, 32, 32, 3)         0         
_________________________________________________________________
create_patches (CreatePatche (None, 8, 8, 48)          0         
_________________________________________________________________
per_patch_fully_connected (P (None, 64, 512)           25088     
_________________________________________________________________
mlp_block (MLPBlock)         (None, 64, 512)           535680    
_________________________________________________________________
mlp_block_1 (MLPBlock)       (None, 64, 512)           535680    
_____________________________________________________________[...]<jupyter_text>Explanation of the layers- The `CreatePatches` layer creates non-overlapping patches of an image using the `tf.image.extract_patches` method. The kernel and stride are provided to ensure the patches are sufficiently small and non-overlapping- The `PerPatchFullyConnected` layer applies a `Dense` layer to all the individual patches of the image, projecting them into `dims` sized vector. This creates our "Patches x Channels" table as described in the MLP-Mixer paper.- The MLPBlock takes no arguements because it maintains the input size. It does the channel mixing.Here we are using some specific hyperparameters:- the internal projection dimension `dims` is set to 512- we are using a depth of 4 MLP-blocks- We are using global average pooling as used in the paper- using guassian dropout with a dropout probablity of 0.9 (due to tendency of model to overfit on dataset of this small size, and even the images are small in size)- Adadelta optimizer with a learning rate of 1.0 (as the optimizer manages lr_decay and gives better results the Adam)<jupyter_code>model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    optimizer=tf.keras.optimizers.Adadelta(1),
    metrics=["accuracy"],
)

b_s = 128 # batch_size

val_dataset = val_dataset.batch(b_s)
history = model.fit(
    x_train,
    y_train,
    validation_data=val_dataset,
    batch_size=b_s,
    epochs=100,
    verbose="auto",
)<jupyter_output>2021-07-21 18:08:37.069792: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
2021-07-21 18:08:37.093193: I tensorflow/core/platform/profile_utils/cpu_utils.cc:114] CPU Frequency: 3800060000 Hz<jupyter_text>As can be seen, the MLP doesnt perform very well for small datasets, and ends up overfitting even with a dropout probablity of 0.9.The best training/val accuracy pair occurs at epoch 84. Train acc: 0.9523, test acc: 0.6132Now, to train a similar model but to use the `Projection` layer instead of Global avg pooling, lets see if that can do any better in this scenario<jupyter_code>history1 = history
def create_model():
    dims = 512
    k, s = (4, 4)  # (kernel, strides)
    depth = 4

    inputs = Input((32, 32, 3))

    x = tf.keras.Sequential(
        [
            RandomFlip(),
            RandomRotation(factor=0.03),
            RandomZoom(height_factor=0.2, width_factor=0.2),
        ],
        name="data_augmentation",
    )(inputs)

    x = CreatePatches(k, s)(x)
    x = PerPatchFullyConnected(dims)(x)

    for _ in range(depth):
        x = MLPBlock()(x)
        x = Permute((2,1))(x)

    # x = Projection()(x)
    x = GlobalAveragePooling1D()(x)

    # classification layer
    x = GaussianDropout(0.9)(x)
    output = Dense(len(classes), activation="softmax", kernel_regularizer="l2")(x)

    return Model(inputs=inputs, outputs=output)


model = create_model()
print(model.summary())
model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(),
    optimizer=tf.keras.optimizers.Adadelta(1),
    metrics=["accuracy"],
)

# b_s = 256 # batch_size

# val_dataset = val_dataset.batch(b_s)
history = model.fit(
    x_train,
    y_train,
    validation_data=val_dataset,
    batch_size=256,
    epochs=100,
    verbose="auto",
)<jupyter_output>2021-07-21 20:50:25.545241: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)
2021-07-21 20:50:25.565189: I tensorflow/core/platform/profile_utils/cpu_utils.cc:114] CPU Frequency: 3800060000 Hz
{'fraction_non_alphanumeric': 0.1044822256568779, 'fraction_numerical': 0.04034003091190108, 'mean_word_length': 3.894856278366112, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '212576', 'license': 'MIT', 'n_tokens_mistral': 2167, 'n_tokens_neox': 1950, 'n_words': 563, 'path': 'src/MLP_Mixer.ipynb', 'repo_name': 'shubhankar000/MLP-Mixer-Tensorflow', 'size': '6470'}