google-research/vision_transformer

This repository, published by Google Research, contains the code and pre-trained models from several research papers on image recognition. The central idea behind the Vision Transformer (ViT) approach is treating an image the same way a language model treats a sequence of words: by slicing the image into small patches and feeding those patches through the same kind of model architecture used in natural language processing. This was a notable departure from how image recognition had traditionally been done, and the papers demonstrate that this approach can match or outperform older methods when trained on large datasets. Alongside the Vision Transformer models, the repository also includes MLP-Mixer, a related architecture that takes a different approach by using only simple matrix operations rather than the attention mechanism. The repository additionally covers follow-up research on how to train these models more effectively, including what data volumes, augmentation techniques, and regularization strategies produce the best results. All the models were pre-trained on large image datasets and are made available for fine-tuning. Fine-tuning means taking one of these pre-trained models and continuing to train it on a smaller, task-specific dataset. The code is written in JAX and Flax, two Python-based frameworks for numerical computing and neural network research developed at Google. The repository includes interactive Jupyter notebooks hosted on Google Colab, which let people experiment with the models through a browser without setting up a local environment. For more serious training runs, the README walks through setting up a cloud-based virtual machine. The full README is longer than what was shown.