Here’s my implementation of MLP-Mixer 20, the all MLP architecture for computer vision without any use of convs and self-attention:

  
      

      github.com
  

  
    

sayakpaul/MLP-Mixer-CIFAR10 85


  Implements MLP-Mixer (https://arxiv.org/abs/2105.01601) with the CIFAR-10 dataset.  

  

  
    
    
  

  


Here’s what is included:

Distributed training with mixed-precision.
Visualization of the token-mixing MLP weights.
A TensorBoard callback to keep track of the learned linear projections of the image patches.

Results are quite competitive with room for improvements for interpretability.