mlp_mixer_6.py

hidden_dims = 128
token_mixing_mlp_dims = 64
channel_mixing_mlp_dims = 128
patch_size = 9
num_classes = 10
num_mixer_layers = 4
input_image_shape = ( 32 , 32 , 3 )

inputs = tf.keras.layers.Input( shape=input_image_shape )

# Conv2D to extract patches
patches = tf.keras.layers.Conv2D( hidden_dims , kernel_size=patch_size , strides=patch_size )( inputs )
# Resizing the patches
patches_reshape = tf.keras.layers.Reshape( ( patches.shape[ 1 ] * patches.shape[ 2 ] , patches.shape[ 3 ] ) )( patches )

x = patches_reshape
for _ in range( num_mixer_layers ):
    x = mixer( x , token_mixing_mlp_dims , channel_mixing_mlp_dims )

# Classifier head
x = tf.keras.layers.LayerNormalization( epsilon=1e-6 )( x )
x = tf.keras.layers.GlobalAveragePooling1D()( x )
outputs = tf.keras.layers.Dense( num_classes , activation='softmax' )( x )

model = tf.keras.models.Model( inputs , outputs )
model.summary()