This is a PyTorch implementation of the paper Patches Are All You Need?.

ConvMixer is Similar to MLP-Mixer. MLP-Mixer separates mixing of spatial and channel dimensions, by applying an MLP across spatial dimension and then an MLP across the channel dimension (spatial MLP replaces the ViT attention and channel MLP is the FFN of ViT).

ConvMixer uses a $1×1$ convolution for channel mixing and a depth-wise convolution for spatial mixing. Since it's a convolution instead of a full MLP across the space, it mixes only the nearby batches in contrast to ViT or MLP-Mixer. Also, the MLP-mixer uses MLPs of two layers for each mixing and ConvMixer uses a single layer for each mixing.

The paper recommends removing the residual connection across the channel mixing (point-wise convolution) and having only a residual connection over the spatial mixing (depth-wise convolution). They also use Batch normalization instead of Layer normalization.

Here's an experiment that trains ConvMixer on CIFAR-10.

```
38import torch
39from torch import nn
40
41from labml_helpers.module import Module
42from labml_nn.utils import clone_module_list
```

`45class ConvMixerLayer(Module):`

`d_model`

is the number of channels in patch embeddings, $h$`kernel_size`

is the size of the kernel of spatial convolution, $k$

`54 def __init__(self, d_model: int, kernel_size: int):`

`59 super().__init__()`

Depth-wise convolution is separate convolution for each channel. We do this with a convolution layer with the number of groups equal to the number of channels. So that each channel is it's own group.

```
63 self.depth_wise_conv = nn.Conv2d(d_model, d_model,
64 kernel_size=kernel_size,
65 groups=d_model,
66 padding=(kernel_size - 1) // 2)
```

Activation after depth-wise convolution

`68 self.act1 = nn.GELU()`

Normalization after depth-wise convolution

`70 self.norm1 = nn.BatchNorm2d(d_model)`

Point-wise convolution is a $1×1$ convolution. i.e. a linear transformation of patch embeddings

`74 self.point_wise_conv = nn.Conv2d(d_model, d_model, kernel_size=1)`

Activation after point-wise convolution

`76 self.act2 = nn.GELU()`

Normalization after point-wise convolution

`78 self.norm2 = nn.BatchNorm2d(d_model)`

`80 def forward(self, x: torch.Tensor):`

For the residual connection around the depth-wise convolution

`82 residual = x`

Depth-wise convolution, activation and normalization

```
85 x = self.depth_wise_conv(x)
86 x = self.act1(x)
87 x = self.norm1(x)
```

Add residual connection

`90 x += residual`

Point-wise convolution, activation and normalization

```
93 x = self.point_wise_conv(x)
94 x = self.act2(x)
95 x = self.norm2(x)
```

`98 return x`

This splits the image into patches of size $p×p$ and gives an embedding for each patch.

`101class PatchEmbeddings(Module):`

`d_model`

is the number of channels in patch embeddings $h$`patch_size`

is the size of the patch, $p$`in_channels`

is the number of channels in the input image (3 for rgb)

`110 def __init__(self, d_model: int, patch_size: int, in_channels: int):`

`116 super().__init__()`

We create a convolution layer with a kernel size and and stride length equal to patch size. This is equivalent to splitting the image into patches and doing a linear transformation on each patch.

`121 self.conv = nn.Conv2d(in_channels, d_model, kernel_size=patch_size, stride=patch_size)`

Activation function

`123 self.act = nn.GELU()`

Batch normalization

`125 self.norm = nn.BatchNorm2d(d_model)`

`x`

is the input image of shape`[batch_size, channels, height, width]`

`127 def forward(self, x: torch.Tensor):`

Apply convolution layer

`132 x = self.conv(x)`

Activation and normalization

```
134 x = self.act(x)
135 x = self.norm(x)
```

`138 return x`

They do average pooling (taking the mean of all patch embeddings) and a final linear transformation to predict the log-probabilities of the image classes.

`141class ClassificationHead(Module):`

`d_model`

is the number of channels in patch embeddings, $h$`n_classes`

is the number of classes in the classification task

`151 def __init__(self, d_model: int, n_classes: int):`

`156 super().__init__()`

Average Pool

`158 self.pool = nn.AdaptiveAvgPool2d((1, 1))`

Linear layer

`160 self.linear = nn.Linear(d_model, n_classes)`

`162 def forward(self, x: torch.Tensor):`

Average pooling

`164 x = self.pool(x)`

Get the embedding, `x`

will have shape `[batch_size, d_model, 1, 1]`

`166 x = x[:, :, 0, 0]`

Linear layer

`168 x = self.linear(x)`

`171 return x`

This combines the patch embeddings block, a number of ConvMixer layers and a classification head.

`174class ConvMixer(Module):`

`conv_mixer_layer`

is a copy of a single ConvMixer layer. We make copies of it to make ConvMixer with`n_layers`

.`n_layers`

is the number of ConvMixer layers (or depth), $d$.`patch_emb`

is the patch embeddings layer.`classification`

is the classification head.

```
181 def __init__(self, conv_mixer_layer: ConvMixerLayer, n_layers: int,
182 patch_emb: PatchEmbeddings,
183 classification: ClassificationHead):
```

`191 super().__init__()`

Patch embeddings

`193 self.patch_emb = patch_emb`

Classification head

`195 self.classification = classification`

Make copies of the ConvMixer layer

`197 self.conv_mixer_layers = clone_module_list(conv_mixer_layer, n_layers)`

`x`

is the input image of shape`[batch_size, channels, height, width]`

`199 def forward(self, x: torch.Tensor):`

Get patch embeddings. This gives a tensor of shape `[batch_size, d_model, height / patch_size, width / patch_size]`

.

`204 x = self.patch_emb(x)`

Pass through ConvMixer layers

```
207 for layer in self.conv_mixer_layers:
208 x = layer(x)
```

Classification head, to get logits

`211 x = self.classification(x)`

`214 return x`