google-research/vision_transformer

Health v2 · maintenance only

Flags: no_releases

Adoption not part of the score

12683 stars · 1476 forks observed · 2026-08-28

What it is AI-extracted, prompt v1, taxonomy v1, 2026-08-29, confidence not recorded

Google Research's official JAX/Flax implementation of Vision Transformer (ViT) and MLP-Mixer architectures, with released pretrained checkpoints on ImageNet and ImageNet-21k. It provides code and Colab notebooks for fine-tuning the models on custom image classification tasks.

Use cases

• fine-tune a pretrained vision transformer on my own image dataset
• run image classification inference with ViT checkpoints
• experiment with MLP-Mixer architectures in JAX
• explore thousands of ViT training checkpoints from the scaling study
• learn how vision transformers are implemented step by step
• zero-shot image classification with LiT models

When to choose

• you want official, well-tested ViT/MLP-Mixer code with pretrained weights
• your stack is JAX/Flax and you need TPU or GPU fine-tuning
• you want to reproduce results from the ViT research papers

When to avoid

• you need PyTorch implementations instead of JAX/Flax
• you need large-scale multi-host pretraining infrastructure (use big_vision)
• you need production image-classification serving rather than research code

Facets

library · maturity stable

machine-learning deep-learning image-processing computer-vision deep-learning machine-learning python cloud vision-transformer jax flax mlp-mixer pretrained-models image-classification fine-tuning research-code gpu

Member repositories

For agents

markdown · JSON · MCP: product_card(name="google-research/vision_transformer")