Vision Transformers in Flax

[conceptofmind]

Hi @marcvanzee ,

I hope all is well.

We had previously messaged about example implementations in Flax. You had told me to open up a Show & Tell Discussion thread about linking an example to the Community Examples page.

I wanted to share a repository with numerous different Vision Transformer architecture implementations and examples in Flax. There are 18 in total: https://github.com/conceptofmind/vit-flax

The repository was co-authored with a peer: https://github.com/lucidrains

The task type would be: Computer Vision

The repository currently includes Flax transformer architecture implementations for papers such as:

• Going deeper with Image Transformers
• CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention
• CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification
• CvT: Introducing Convolutions to Vision Transformers
• DeepViT: Towards Deeper Vision Transformer
• LeViT: a Vision Transformer in ConvNet's Clothing for Faster Inference
• Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding
• Three things everyone should know about Vision Transformers
• Learning to Merge Tokens in Vision Transformers
• Rethinking Spatial Dimensions of Vision Transformers
• RegionViT: Regional-to-Local Attention for Vision Transformers
• ScalableViT: Rethinking the Context-oriented Generalization of Vision Transformer
• Better plain ViT baselines for ImageNet-1k
• Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet
• Twins: Revisiting the Design of Spatial Attention in Vision Transformers
• Vision Transformer for Small-Size Datasets

I will be continuously updating it as I finish more research paper implementations.

I greatly appreciate your time and consideration.

Thank you,

Enrico Shippole

[cgarciae]

Hey @conceptofmind, thanks for sharing! I thinks its important to point out that Flax implementations of the Vision Transformer are also available in the big_vision and vision_transformer repos, code from the original authors if I am not mistaken so its also worth to check out.

[conceptofmind]

Hi @cgarciae ,

I appreciate the additional information.

I had not used any official Google research code as a reference when working on the different architectures. So, it is very possible that I may have reimplemented one of the Vision Transformers already available in big_vision and vision_transformer in Flax without realizing it. Would this be an issue at all? I will make any necessary changes.

I will have to review the source files in those repositories you linked to see if there are any other official Google Flax implementations such as the original Vision Transformer one.

I did make sure to include links for all official Github repositories, research papers, and bibtex citations for each model in the example vit-flax repository.

Thank you for your help.

Enrico Shippole

[cgarciae]

Hey @conceptofmind! No issue at all 🙂 I just wanted to highlight that work as part of the conversation since it just happens to be one of the examples where the original research was written in Flax.

I think your repo could be show-cased in the Community Examples section, feel free so send a PR.

[conceptofmind]

Hi @cgarciae ,

I appreciate it. I will send a PR over now.

Thank you very much!

Enrico Shippole