MLP-Mixer in Flax and PyTorch

In this video, we implement the MLP-mixer in both Flax and PyTorch. It is a recent model for image classification that only uses simple multilayer perceptron blocks, however, it seems to perform as well as CNNs and Vision Transformers. Conceptually, it is really simple and the implementation is straightforward and therefore we try to code it up in two different deep learning frameworks - PyTorch and Flax. If you are new to Flax do not worry because this video contains a quick tutorial on the most important concepts. Last but not least we also investigate the relationship between the MLP-Mixer and Convolutional neural networks.

Paper: https://arxiv.org/abs/2105.01601
Official implementation: https://github.com/google-research/vi... (check out the branch linen if necessary)
Video implementation: https://github.com/jankrepl/mildlyove...

00:00 Intro
01:12 High level explanation
03:13 Flax 101
07:39 Flax implementation
08:48 nn.Dense behavior
10:11 Flax implementation continued
12:43 Torch: MlpBlock
14:14 Torch: MixerBlock
17:45 Channel mixing as convolution
20:06 Token mixing as convolution
22:43 Torch: MlpMixer
25:58 Patch embedding without convolution
28:27 Torch: MlpMixer continued
28:56 Comparing Flax and Torch networks
31:50 Outro


Wanna learn more about the LayerNorm?    • Vision Transformer in PyTorch  

If you have any video suggestions or you just wanna chat feel free to join the discord server:   / discord  

Credits logo animation
Title: Conjungation · Author: Uncle Milk · Source:   / unclemilk   · License: https://creativecommons.org/licenses/... · Download (9MB): https://auboutdufil.com/?id=600