{"id":13440142,"url":"https://github.com/google-research/vision_transformer","last_synced_at":"2025-05-16T12:00:17.323Z","repository":{"id":36996853,"uuid":"306018946","full_name":"google-research/vision_transformer","owner":"google-research","description":null,"archived":false,"fork":false,"pushed_at":"2025-03-06T03:14:39.000Z","size":145223,"stargazers_count":11336,"open_issues_count":134,"forks_count":1376,"subscribers_count":108,"default_branch":"main","last_synced_at":"2025-05-09T11:25:36.944Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/google-research.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2020-10-21T12:35:02.000Z","updated_at":"2025-05-09T09:04:16.000Z","dependencies_parsed_at":"2024-01-06T04:16:42.498Z","dependency_job_id":"a8056c75-3d54-4247-ae56-d29bbc06cd69","html_url":"https://github.com/google-research/vision_transformer","commit_stats":{"total_commits":174,"total_committers":26,"mean_commits":"6.6923076923076925","dds":0.3275862068965517,"last_synced_commit":"2ecbc73207056d8fcbd7513ab43cd417a0b385b3"},"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-research%2Fvision_transformer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-research%2Fvision_transformer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-research%2Fvision_transformer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/google-research%2Fvision_transformer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/google-research","download_url":"https://codeload.github.com/google-research/vision_transformer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254527069,"owners_count":22085917,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-07-31T03:01:20.141Z","updated_at":"2025-05-16T12:00:17.275Z","avatar_url":"https://github.com/google-research.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook","DLA","Models","Vision Model Backbone","其他_机器视觉","JAX Models","Table of Contents","Fundamental MIM Methods","Models and Projects"],"sub_categories":["ViT","网络服务_其他","Reinforcement Learning","微软Transformer霸榜模型","MIM for Transformers","Flax"],"readme":"# Vision Transformer and MLP-Mixer Architectures\n\nIn this repository we release models from the papers\n\n- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929)\n- [MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601)\n- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://arxiv.org/abs/2106.10270)\n- [When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations](https://arxiv.org/abs/2106.01548)\n- [LiT: Zero-Shot Transfer with Locked-image text Tuning](https://arxiv.org/abs/2111.07991)\n- [Surrogate Gap Minimization Improves Sharpness-Aware Training](https://arxiv.org/abs/2203.08065)\n\nThe models were pre-trained on the [ImageNet](http://www.image-net.org/) and\n[ImageNet-21k](http://www.image-net.org/) datasets. We provide the code for\nfine-tuning the released models in\n[JAX](https://jax.readthedocs.io)/[Flax](http://flax.readthedocs.io).\n\nThe models from this codebase were originally trained in\nhttps://github.com/google-research/big_vision/\nwhere you can find more advanced code (e.g. multi-host training), as well as\nsome of the original training scripts (e.g.\n[configs/vit_i21k.py](https://github.com/google-research/big_vision/blob/main/big_vision/configs/vit_i21k.py)\nfor pre-training a ViT, or\n[configs/transfer.py](https://github.com/google-research/big_vision/blob/main/big_vision/configs/transfer.py)\nfor transfering a model).\n\nTable of contents:\n\n- [Vision Transformer and MLP-Mixer Architectures](#vision-transformer-and-mlp-mixer-architectures)\n\t- [Colab](#colab)\n\t- [Installation](#installation)\n\t- [Fine-tuning a model](#fine-tuning-a-model)\n\t- [Vision Transformer](#vision-transformer)\n\t\t- [Available ViT models](#available-vit-models)\n\t\t- [Expected ViT results](#expected-vit-results)\n\t- [MLP-Mixer](#mlp-mixer)\n\t\t- [Available Mixer models](#available-mixer-models)\n\t\t- [Expected Mixer results](#expected-mixer-results)\n\t- [LiT models](#lit-models)\n\t- [Running on cloud](#running-on-cloud)\n\t\t- [Create a VM](#create-a-vm)\n\t\t- [Setup VM](#setup-vm)\n\t- [Bibtex](#bibtex)\n\t- [Disclaimers](#disclaimers)\n\t- [Changelog](#changelog)\n\n\n## Colab\n\nBelow Colabs run both with GPUs, and TPUs (8 cores, data parallelism).\n\nThe first Colab demonstrates the JAX code of Vision Transformers and MLP Mixers.\nThis Colab allows you to edit the files from the repository directly in the\nColab UI and has annotated Colab cells that walk you through the code step by\nstep, and lets you interact with the data.\n\nhttps://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax.ipynb\n\nThe second Colab allows you to explore the \u003e50k Vision Transformer and hybrid\ncheckpoints that were used to generate the data of the third paper \"How to train\nyour ViT? ...\". The Colab includes code to explore and select checkpoints, and\nto do inference both using the JAX code from this repo, and also using the\npopular [`timm`] PyTorch library that can directly load these checkpoints as\nwell. Note that a handful of models are also available directly from TF-Hub:\n[sayakpaul/collections/vision_transformer] (external contribution by [Sayak\nPaul]).\n\nThe second Colab also lets you fine-tune the checkpoints on any tfds dataset\nand your own dataset with examples in individual JPEG files (optionally directly\nreading from Google Drive).\n\nhttps://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax_augreg.ipynb\n\n**Note**: As for now (6/20/21) Google Colab only supports a single GPU (Nvidia\nTesla T4), and TPUs (currently TPUv2-8) are attached indirectly to the Colab VM\nand communicate over slow network, which leads to pretty bad training speed. You\nwould usually want to set up a dedicated machine if you have a non-trivial\namount of data to fine-tune on. For details see the\n[Running on cloud](#running-on-cloud) section.\n\n\n[`timm`]: https://github.com/rwightman/pytorch-image-models\n[sayakpaul/collections/vision_transformer]: https://tfhub.dev/sayakpaul/collections/vision_transformer\n[Sayak Paul]: https://github.com/sayakpaul\n\n\n\n## Installation\n\nMake sure you have `Python\u003e=3.10` installed on your machine.\n\nInstall JAX and python dependencies by running:\n\n```\n# If using GPU:\npip install -r vit_jax/requirements.txt\n\n# If using TPU:\npip install -r vit_jax/requirements-tpu.txt\n```\n\nFor newer versions of [JAX](https://github.com/google/jax), follow the instructions\nprovided in the corresponding repository linked here. Note that installation\ninstructions for CPU, GPU and TPU differs slightly.\n\nInstall [Flaxformer](https://github.com/google/flaxformer), follow the instructions\nprovided in the corresponding repository linked here.\n\nFor more details refer to the section [Running on cloud](#running-on-cloud)\nbelow.\n\n\n## Fine-tuning a model\n\nYou can run fine-tuning of the downloaded model on your dataset of interest. All\nmodels share the same command line interface.\n\nFor example for fine-tuning a ViT-B/16 (pre-trained on imagenet21k) on CIFAR10\n(note how we specify `b16,cifar10` as arguments to the config, and how we\ninstruct the code to access the models directly from a GCS bucket instead of\nfirst downloading them into the local directory):\n\n```bash\npython -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \\\n    --config=$(pwd)/vit_jax/configs/vit.py:b16,cifar10 \\\n    --config.pretrained_dir='gs://vit_models/imagenet21k'\n```\n\nIn order to fine-tune a Mixer-B/16 (pre-trained on imagenet21k) on CIFAR10:\n\n```bash\npython -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \\\n    --config=$(pwd)/vit_jax/configs/mixer_base16_cifar10.py \\\n    --config.pretrained_dir='gs://mixer_models/imagenet21k'\n```\n\nThe \"How to train your ViT? ...\" paper added \u003e50k checkpoints that you can\nfine-tune with the [`configs/augreg.py`] config. When you only specify the model\nname (the `config.name` value from [`configs/model.py`]), then the best i21k\ncheckpoint by upstream validation accuracy (\"recommended\" checkpoint, see\nsection 4.5 of the paper) is chosen. To make up your mind which model you want\nto use, have a look at Figure 3 in the paper. It's also possible to choose a\ndifferent checkpoint (see Colab [`vit_jax_augreg.ipynb`]) and then specify the\nvalue from the `filename` or `adapt_filename` column, which correspond to the\nfilenames without `.npz` from the [`gs://vit_models/augreg`] directory.\n\n```bash\npython -m vit_jax.main --workdir=/tmp/vit-$(date +%s) \\\n    --config=$(pwd)/vit_jax/configs/augreg.py:R_Ti_16 \\\n    --config.dataset=oxford_iiit_pet \\\n    --config.base_lr=0.01\n```\n\nCurrently, the code will automatically download CIFAR-10 and CIFAR-100 datasets.\nOther public or custom datasets can be easily integrated, using [tensorflow\ndatasets library](https://github.com/tensorflow/datasets/). Note that you will\nalso need to update `vit_jax/input_pipeline.py` to specify some parameters about\nany added dataset.\n\nNote that our code uses all available GPUs/TPUs for fine-tuning.\n\nTo see a detailed list of all available flags, run `python3 -m vit_jax.train\n--help`.\n\nNotes on memory:\n\n- Different models require different amount of memory. Available memory also\n  depends on the accelerator configuration (both type and count). If you\n  encounter an out-of-memory error you can increase the value of\n  `--config.accum_steps=8` -- alternatively, you could also decrease the\n  `--config.batch=512` (and decrease `--config.base_lr` accordingly).\n- The host keeps a shuffle buffer in memory. If you encounter a host OOM (as\n  opposed to an accelerator OOM), you can decrease the default\n  `--config.shuffle_buffer=50000`.\n\n\n## Vision Transformer\n\nby Alexey Dosovitskiy\\*†, Lucas Beyer\\*, Alexander Kolesnikov\\*, Dirk\nWeissenborn\\*, Xiaohua Zhai\\*, Thomas Unterthiner, Mostafa Dehghani, Matthias\nMinderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit and Neil Houlsby\\*†.\n\n(\\*) equal technical contribution, (†) equal advising.\n\n![Figure 1 from paper](vit_figure.png)\n\nOverview of the model: we split an image into fixed-size patches, linearly embed\neach of them, add position embeddings, and feed the resulting sequence of\nvectors to a standard Transformer encoder. In order to perform classification,\nwe use the standard approach of adding an extra learnable \"classification token\"\nto the sequence.\n\n### Available ViT models\n\nWe provide a variety of ViT models in different GCS buckets. The models can be\ndownloaded with e.g.:\n\n```\nwget https://storage.googleapis.com/vit_models/imagenet21k/ViT-B_16.npz\n```\n\nThe model filenames (without the `.npz` extension) correspond to the\n`config.model_name` in [`vit_jax/configs/models.py`]\n\n- [`gs://vit_models/imagenet21k`] - Models pre-trained on ImageNet-21k.\n- [`gs://vit_models/imagenet21k+imagenet2012`] - Models pre-trained on\n  ImageNet-21k and fine-tuned on ImageNet.\n- [`gs://vit_models/augreg`] - Models pre-trained on ImageNet-21k,\n  applying varying amounts of [AugReg]. Improved performance.\n- [`gs://vit_models/sam`] - Models pre-trained on ImageNet with [SAM].\n- [`gs://vit_models/gsam`] - Models pre-trained on ImageNet with [GSAM].\n\nWe recommend using the following checkpoints, trained with [AugReg] that have\nthe best pre-training metrics:\n\n|  Model   |                                   Pre-trained checkpoint                                   |   Size   |                                                       Fine-tuned checkpoint                                                        | Resolution | Img/sec | Imagenet accuracy |\n| :------- | :----------------------------------------------------------------------------------------- | -------: | :--------------------------------------------------------------------------------------------------------------------------------- | ---------: | ------: | ----------------: |\n| L/16     | `gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0.npz`     | 1243 MiB | `gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_strong1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz`     |        384 |      50 |            85.59% |\n| B/16     | `gs://vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0.npz`     |  391 MiB | `gs://vit_models/augreg/B_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz`     |        384 |     138 |            85.49% |\n| S/16     | `gs://vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0.npz`     |  115 MiB | `gs://vit_models/augreg/S_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz`     |        384 |     300 |            83.73% |\n| R50+L/32 | `gs://vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1.npz` | 1337 MiB | `gs://vit_models/augreg/R50_L_32-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz` |        384 |     327 |            85.99% |\n| R26+S/32 | `gs://vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0.npz`  |  170 MiB | `gs://vit_models/augreg/R26_S_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz`  |        384 |     560 |            83.85% |\n| Ti/16    | `gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz`      |   37 MiB | `gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz`      |        384 |     610 |            78.22% |\n| B/32     | `gs://vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0.npz`      |  398 MiB | `gs://vit_models/augreg/B_32-i21k-300ep-lr_0.001-aug_light1-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz`      |        384 |     955 |            83.59% |\n| S/32     | `gs://vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0.npz`        |  118 MiB | `gs://vit_models/augreg/S_32-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.01-res_384.npz`        |        384 |    2154 |            79.58% |\n| R+Ti/16  | `gs://vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz`    |   40 MiB | `gs://vit_models/augreg/R_Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_384.npz`    |        384 |    2426 |            75.40% |\n\nThe results from the original ViT paper (https://arxiv.org/abs/2010.11929) have\nbeen replicated using the models from [`gs://vit_models/imagenet21k`]:\n\n| model        | dataset      | dropout=0.0                                                                                                                                                         | dropout=0.1                                                                                                                                                          |\n|:-------------|:-------------|:--------------------------------------------------------------------------------------------------------------------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| R50+ViT-B_16 | cifar10      | 98.72%, 3.9h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/cifar10/do_0.0\u0026_smoothingWeight=0)      | 98.94%, 10.1h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/cifar10/do_0.1\u0026_smoothingWeight=0)      |\n| R50+ViT-B_16 | cifar100     | 90.88%, 4.1h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/cifar100/do_0.0\u0026_smoothingWeight=0)     | 92.30%, 10.1h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/cifar100/do_0.1\u0026_smoothingWeight=0)     |\n| R50+ViT-B_16 | imagenet2012 | 83.72%, 9.9h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/imagenet2012/do_0.0\u0026_smoothingWeight=0) | 85.08%, 24.2h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5ER50.ViT-B_16/imagenet2012/do_0.1\u0026_smoothingWeight=0) |\n| ViT-B_16     | cifar10      | 99.02%, 2.2h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/cifar10/do_0.0\u0026_smoothingWeight=0)          | 98.76%, 7.8h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/cifar10/do_0.1\u0026_smoothingWeight=0)           |\n| ViT-B_16     | cifar100     | 92.06%, 2.2h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/cifar100/do_0.0\u0026_smoothingWeight=0)         | 91.92%, 7.8h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/cifar100/do_0.1\u0026_smoothingWeight=0)          |\n| ViT-B_16     | imagenet2012 | 84.53%, 6.5h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/imagenet2012/do_0.0\u0026_smoothingWeight=0)     | 84.12%, 19.3h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_16/imagenet2012/do_0.1\u0026_smoothingWeight=0)     |\n| ViT-B_32     | cifar10      | 98.88%, 0.8h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/cifar10/do_0.0\u0026_smoothingWeight=0)          | 98.75%, 1.8h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/cifar10/do_0.1\u0026_smoothingWeight=0)           |\n| ViT-B_32     | cifar100     | 92.31%, 0.8h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/cifar100/do_0.0\u0026_smoothingWeight=0)         | 92.05%, 1.8h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/cifar100/do_0.1\u0026_smoothingWeight=0)          |\n| ViT-B_32     | imagenet2012 | 81.66%, 3.3h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/imagenet2012/do_0.0\u0026_smoothingWeight=0)     | 81.31%, 4.9h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-B_32/imagenet2012/do_0.1\u0026_smoothingWeight=0)      |\n| ViT-L_16     | cifar10      | 99.13%, 6.9h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/cifar10/do_0.0\u0026_smoothingWeight=0)          | 99.14%, 24.7h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/cifar10/do_0.1\u0026_smoothingWeight=0)          |\n| ViT-L_16     | cifar100     | 92.91%, 7.1h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/cifar100/do_0.0\u0026_smoothingWeight=0)         | 93.22%, 24.4h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/cifar100/do_0.1\u0026_smoothingWeight=0)         |\n| ViT-L_16     | imagenet2012 | 84.47%, 16.8h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/imagenet2012/do_0.0\u0026_smoothingWeight=0)    | 85.05%, 59.7h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_16/imagenet2012/do_0.1\u0026_smoothingWeight=0)     |\n| ViT-L_32     | cifar10      | 99.06%, 1.9h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/cifar10/do_0.0\u0026_smoothingWeight=0)          | 99.09%, 6.1h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/cifar10/do_0.1\u0026_smoothingWeight=0)           |\n| ViT-L_32     | cifar100     | 93.29%, 1.9h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/cifar100/do_0.0\u0026_smoothingWeight=0)         | 93.34%, 6.2h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/cifar100/do_0.1\u0026_smoothingWeight=0)          |\n| ViT-L_32     | imagenet2012 | 81.89%, 7.5h (A100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/imagenet2012/do_0.0\u0026_smoothingWeight=0)     | 81.13%, 15.0h (V100), [tb.dev](https://tensorboard.dev/experiment/nwXQNjudRJW3dtQzhPZwwA/#scalars\u0026regexInput=%5EViT-L_32/imagenet2012/do_0.1\u0026_smoothingWeight=0)     |\n\nWe also would like to emphasize that high-quality results can be achieved with\nshorter training schedules and encourage users of our code to play with\nhyper-parameters to trade-off accuracy and computational budget.\nSome examples for CIFAR-10/100 datasets are presented in the table below.\n\n| upstream    | model    | dataset      | total_steps / warmup_steps  | accuracy | wall-clock time |                                                                         link |\n| ----------- | -------- | ------------ | --------------------------- | -------- | --------------- | ---------------------------------------------------------------------------- |\n| imagenet21k | ViT-B_16 | cifar10      | 500 / 50                    |   98.59% |             17m | [tensorboard.dev](https://tensorboard.dev/experiment/QgkpiW53RPmjkabe1ME31g/) |\n| imagenet21k | ViT-B_16 | cifar10      | 1000 / 100                  |   98.86% |             39m | [tensorboard.dev](https://tensorboard.dev/experiment/w8DQkDeJTOqJW5js80gOQg/) |\n| imagenet21k | ViT-B_16 | cifar100     | 500 / 50                    |   89.17% |             17m | [tensorboard.dev](https://tensorboard.dev/experiment/5hM4GrnAR0KEZg725Ewnqg/) |\n| imagenet21k | ViT-B_16 | cifar100     | 1000 / 100                  |   91.15% |             39m | [tensorboard.dev](https://tensorboard.dev/experiment/QLQTaaIoT9uEcAjtA0eRwg/) |\n\n\n## MLP-Mixer\n\nby Ilya Tolstikhin\\*, Neil Houlsby\\*, Alexander Kolesnikov\\*, Lucas Beyer\\*,\nXiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers,\nJakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy.\n\n(\\*) equal contribution.\n\n![Figure 1 from paper](mixer_figure.png)\n\nMLP-Mixer (*Mixer* for short) consists of per-patch linear embeddings, Mixer\nlayers, and a classifier head. Mixer layers contain one token-mixing MLP and one\nchannel-mixing MLP, each consisting of two fully-connected layers and a GELU\nnonlinearity. Other components include: skip-connections, dropout, and linear\nclassifier head.\n\nFor installation follow [the same steps](#installation) as above.\n\n### Available Mixer models\n\nWe provide the Mixer-B/16 and Mixer-L/16 models pre-trained on the ImageNet and\nImageNet-21k datasets. Details can be found in Table 3 of the Mixer paper. All\nthe models can be found at:\n\nhttps://console.cloud.google.com/storage/mixer_models/\n\nNote that these models are also available directly from TF-Hub:\n[sayakpaul/collections/mlp-mixer] (external contribution by [Sayak\nPaul]).\n\n[sayakpaul/collections/mlp-mixer]: https://tfhub.dev/sayakpaul/collections/mlp-mixer\n\n### Expected Mixer results\n\nWe ran the fine-tuning code on Google Cloud machine with four V100 GPUs with the\ndefault adaption parameters from this repository. Here are the results:\n\nupstream     | model      | dataset | accuracy | wall_clock_time | link\n:----------- | :--------- | :------ | -------: | :-------------- | :---\nImageNet     | Mixer-B/16 | cifar10 | 96.72%   | 3.0h            | [tensorboard.dev](https://tensorboard.dev/experiment/j9zCYt9yQVm93nqnsDZayA/)\nImageNet     | Mixer-L/16 | cifar10 | 96.59%   | 3.0h            | [tensorboard.dev](https://tensorboard.dev/experiment/Q4feeErzRGGop5XzAvYj2g/)\nImageNet-21k | Mixer-B/16 | cifar10 | 96.82%   | 9.6h            | [tensorboard.dev](https://tensorboard.dev/experiment/mvP4McV2SEGFeIww20ie5Q/)\nImageNet-21k | Mixer-L/16 | cifar10 | 98.34%   | 10.0h           | [tensorboard.dev](https://tensorboard.dev/experiment/dolAJyQYTYmudytjalF6Jg/)\n\n\n## LiT models\n\nFor details, refer to the Google AI blog post\n[LiT: adding language understanding to image models](http://ai.googleblog.com/2022/04/locked-image-tuning-adding-language.html),\nor read the CVPR paper \"LiT: Zero-Shot Transfer with Locked-image text Tuning\"\n(https://arxiv.org/abs/2111.07991).\n\nWe published a Transformer B/16-base model with an ImageNet zeroshot accuracy of\n72.1%, and a L/16-large model with an ImageNet zeroshot accuracy of 75.7%. For\nmore details about these models, please refer to the\n[LiT model card](model_cards/lit.md).\n\nWe provide a in-browser demo with small text encoders for interactive use (the\nsmallest models should even run on a modern cell phone):\n\nhttps://google-research.github.io/vision_transformer/lit/\n\nAnd finally a Colab to use the JAX models with both image and text encoders:\n\nhttps://colab.research.google.com/github/google-research/vision_transformer/blob/main/lit.ipynb\n\nNote that none of above models support multi-lingual inputs yet, but we're\nworking on publishing such models and will update this repository once they\nbecome available.\n\nThis repository only contains evaluation code for LiT models. You can find the\ntraining code in the `big_vision` repository:\n\nhttps://github.com/google-research/big_vision/tree/main/big_vision/configs/proj/image_text\n\nExpected zeroshot results from [`model_cards/lit.md`] (note that the zeroshot\nevaluation is slightly different from the simplified evaluation in the Colab):\n\n| Model | B16B_2 | L16L |\n| :--- | ---: | ---: |\n| ImageNet zero-shot | 73.9% | 75.7% |\n| ImageNet v2 zero-shot | 65.1% | 66.6% |\n| CIFAR100 zero-shot | 79.0% | 80.5% |\n| Pets37 zero-shot | 83.3% | 83.3% |\n| Resisc45 zero-shot | 25.3% | 25.6% |\n| MS-COCO Captions image-to-text retrieval | 51.6% | 48.5% |\n| MS-COCO Captions text-to-image retrieval | 31.8% | 31.1% |\n\n## Running on cloud\n\nWhile above [colabs](#colab) are pretty useful to get started, you would usually\nwant to train on a larger machine with more powerful accelerators.\n\n### Create a VM\n\nYou can use the following commands to setup a VM with GPUs on Google Cloud:\n\n```bash\n# Set variables used by all commands below.\n# Note that project must have accounting set up.\n# For a list of zones with GPUs refer to\n# https://cloud.google.com/compute/docs/gpus/gpu-regions-zones\nPROJECT=my-awesome-gcp-project  # Project must have billing enabled.\nVM_NAME=vit-jax-vm-gpu\nZONE=europe-west4-b\n\n# Below settings have been tested with this repository. You can choose other\n# combinations of images \u0026 machines (e.g.), refer to the corresponding gcloud commands:\n# gcloud compute images list --project ml-images\n# gcloud compute machine-types list\n# etc.\ngcloud compute instances create $VM_NAME \\\n    --project=$PROJECT --zone=$ZONE \\\n    --image=c1-deeplearning-tf-2-5-cu110-v20210527-debian-10 \\\n    --image-project=ml-images --machine-type=n1-standard-96 \\\n    --scopes=cloud-platform,storage-full --boot-disk-size=256GB \\\n    --boot-disk-type=pd-ssd --metadata=install-nvidia-driver=True \\\n    --maintenance-policy=TERMINATE \\\n    --accelerator=type=nvidia-tesla-v100,count=8\n\n# Connect to VM (after some minutes needed to setup \u0026 start the machine).\ngcloud compute ssh --project $PROJECT --zone $ZONE $VM_NAME\n\n# Stop the VM after use (only storage is billed for a stopped VM).\ngcloud compute instances stop --project $PROJECT --zone $ZONE $VM_NAME\n\n# Delete VM after use (this will also remove all data stored on VM).\ngcloud compute instances delete --project $PROJECT --zone $ZONE $VM_NAME\n```\n\nAlternatively, you can use the following similar commands to set up a Cloud VM\nwith TPUs attached to them (below commands copied from the [TPU tutorial]):\n\n[TPU tutorial]: https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm\n\n```bash\nPROJECT=my-awesome-gcp-project  # Project must have billing enabled.\nVM_NAME=vit-jax-vm-tpu\nZONE=europe-west4-a\n\n# Required to set up service identity initially.\ngcloud beta services identity create --service tpu.googleapis.com\n\n# Create a VM with TPUs directly attached to it.\ngcloud alpha compute tpus tpu-vm create $VM_NAME \\\n    --project=$PROJECT --zone=$ZONE \\\n    --accelerator-type v3-8 \\\n    --version tpu-vm-base\n\n# Connect to VM (after some minutes needed to setup \u0026 start the machine).\ngcloud alpha compute tpus tpu-vm ssh --project $PROJECT --zone $ZONE $VM_NAME\n\n# Stop the VM after use (only storage is billed for a stopped VM).\ngcloud alpha compute tpus tpu-vm stop --project $PROJECT --zone $ZONE $VM_NAME\n\n# Delete VM after use (this will also remove all data stored on VM).\ngcloud alpha compute tpus tpu-vm delete --project $PROJECT --zone $ZONE $VM_NAME\n```\n\n### Setup VM\n\nAnd then fetch the repository and the install dependencies (including `jaxlib`\nwith TPU support) as usual:\n\n```bash\ngit clone --depth=1 --branch=master https://github.com/google-research/vision_transformer\ncd vision_transformer\n\n# optional: install virtualenv\npip3 install virtualenv\npython3 -m virtualenv env\n. env/bin/activate\n```\n\nIf you're connected to a VM with GPUs attached, install JAX and other dependencies with the following\ncommand:\n\n```bash\npip install -r vit_jax/requirements.txt\n```\n\nIf you're connected to a VM with TPUs attached, install JAX and other dependencies with the following\ncommand:\n\n```bash\npip install -r vit_jax/requirements-tpu.txt\n```\n\nInstall [Flaxformer](https://github.com/google/flaxformer), follow the instructions\nprovided in the corresponding repository linked here.\n\nFor both GPUs and TPUs, Check that JAX can connect to attached accelerators with the command:\n```bash\npython -c 'import jax; print(jax.devices())'\n```\n\nAnd finally execute one of the commands mentioned in the section\n[fine-tuning a model](#fine-tuning-a-model).\n\n\n## Bibtex\n\n```\n@article{dosovitskiy2020vit,\n  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},\n  author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and  Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and Uszkoreit, Jakob and Houlsby, Neil},\n  journal={ICLR},\n  year={2021}\n}\n\n@article{tolstikhin2021mixer,\n  title={MLP-Mixer: An all-MLP Architecture for Vision},\n  author={Tolstikhin, Ilya and Houlsby, Neil and Kolesnikov, Alexander and Beyer, Lucas and Zhai, Xiaohua and Unterthiner, Thomas and Yung, Jessica and Steiner, Andreas and Keysers, Daniel and Uszkoreit, Jakob and Lucic, Mario and Dosovitskiy, Alexey},\n  journal={arXiv preprint arXiv:2105.01601},\n  year={2021}\n}\n\n@article{steiner2021augreg,\n  title={How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers},\n  author={Steiner, Andreas and Kolesnikov, Alexander and and Zhai, Xiaohua and Wightman, Ross and Uszkoreit, Jakob and Beyer, Lucas},\n  journal={arXiv preprint arXiv:2106.10270},\n  year={2021}\n}\n\n@article{chen2021outperform,\n  title={When Vision Transformers Outperform ResNets without Pretraining or Strong Data Augmentations},\n  author={Chen, Xiangning and Hsieh, Cho-Jui and Gong, Boqing},\n  journal={arXiv preprint arXiv:2106.01548},\n  year={2021},\n}\n\n@article{zhuang2022gsam,\n  title={Surrogate Gap Minimization Improves Sharpness-Aware Training},\n  author={Zhuang, Juntang and Gong, Boqing and Yuan, Liangzhe and Cui, Yin and Adam, Hartwig and Dvornek, Nicha and Tatikonda, Sekhar and Duncan, James and Liu, Ting},\n  journal={ICLR},\n  year={2022},\n}\n\n@article{zhai2022lit,\n  title={LiT: Zero-Shot Transfer with Locked-image Text Tuning},\n  author={Zhai, Xiaohua and Wang, Xiao and Mustafa, Basil and Steiner, Andreas and Keysers, Daniel and Kolesnikov, Alexander and Beyer, Lucas},\n  journal={CVPR},\n  year={2022}\n}\n```\n\n\n## Changelog\n\nIn reverse chronological order:\n\n- 2022-08-18: Added LiT-B16B_2 model that was trained for 60k steps\n  (LiT_B16B: 30k) without linear head on the image side (LiT_B16B: 768) and has\n  better performance.\n\n- 2022-06-09: Added the ViT and Mixer models trained from scratch using\n  [GSAM] on ImageNet without strong data augmentations. The resultant ViTs\n  outperform those of similar sizes trained using AdamW optimizer or the\n  original [SAM] algorithm, or with strong data augmentations.\n\n- 2022-04-14: Added models and Colab for [LiT models](#lit-models).\n\n- 2021-07-29: Added ViT-B/8 AugReg models (3 upstream checkpoints and adaptations\n  with resolution=224).\n\n- 2021-07-02: Added the \"When Vision Transformers Outperform\n  ResNets...\" paper\n\n- 2021-07-02: Added [SAM](https://arxiv.org/abs/2010.01412)\n  (Sharpness-Aware Minimization) optimized ViT and MLP-Mixer checkpoints.\n\n- 2021-06-20: Added the \"How to train your ViT? ...\" paper, and a new\n  Colab to explore the \u003e50k pre-trained and fine-tuned checkpoints mentioned in\n  the paper.\n\n- 2021-06-18: This repository was rewritten to use Flax Linen API and\n  `ml_collections.ConfigDict` for configuration.\n\n- 2021-05-19: With publication of the \"How to train your ViT? ...\"\n  paper, we added more than 50k ViT and hybrid models pre-trained on ImageNet and\n  ImageNet-21k with various degrees of data augmentation and model regularization,\n  and fine-tuned on ImageNet, Pets37, Kitti-distance, CIFAR-100, and Resisc45.\n  Check out [`vit_jax_augreg.ipynb`] to navigate this treasure trove of models!\n  For example, you can use that Colab to fetch the filenames of recommended\n  pre-trained and fine-tuned checkpoints from the `i21k_300` column of Table 3 in\n  the paper.\n\n- 2020-12-01: Added the R50+ViT-B/16 hybrid model (ViT-B/16 on\n  top of a Resnet-50 backbone). When pretrained on imagenet21k, this model\n  achieves almost the performance of the L/16 model with less than half the\n  computational finetuning cost. Note that \"R50\" is somewhat modified for the\n  B/16 variant: The original ResNet-50 has [3,4,6,3] blocks, each reducing the\n  resolution of the image by a factor of two. In combination with the ResNet\n  stem this would result in a reduction of 32x so even with a patch size of\n  (1,1) the ViT-B/16 variant cannot be realized anymore. For this reason we\n  instead use [3,4,9] blocks for the R50+B/16 variant.\n\n- 2020-11-09: Added the ViT-L/16 model.\n\n- 2020-10-29: Added ViT-B/16 and ViT-L/16 models pretrained\n  on ImageNet-21k and then fine-tuned on ImageNet at 224x224 resolution (instead\n  of default 384x384). These models have the suffix \"-224\" in their name.\n  They are expected to achieve 81.2% and 82.7% top-1 accuracies respectively.\n\n\n## Disclaimers\n\nOpen source release prepared by Andreas Steiner.\n\nNote: This repository was forked and modified from\n[google-research/big_transfer](https://github.com/google-research/big_transfer).\n\n**This is not an official Google product.**\n\n\n[GSAM]: https://arxiv.org/abs/2203.08065\n[SAM]: https://arxiv.org/abs/2010.01412\n[AugReg]: https://arxiv.org/abs/2106.10270\n\n[`vit_jax/configs/models.py`]: https://github.com/google-research/vision_transformer/blob/main/vit_jax/configs/models.py\n[`model_cards/lit.md`]: https://github.com/google-research/vision_transformer/blob/main/model_cards/lit.md\n\n[`configs/augreg.py`]: https://github.com/google-research/vision_transformer/blob/main/vit_jax/configs/augreg.py\n[`configs/model.py`]: https://github.com/google-research/vision_transformer/blob/main/vit_jax/configs/models.py\n[`vit_jax_augreg.ipynb`]: https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax_augreg.ipynb\n[`vit_jax.ipynb`]: https://colab.research.google.com/github/google-research/vision_transformer/blob/main/vit_jax.ipynb\n\n[`gs://vit_models/imagenet21k`]: https://console.cloud.google.com/storage/browser/vit_models/imagenet21k/\n[`gs://vit_models/imagenet21k+imagenet2012`]: https://console.cloud.google.com/storage/browser/vit_models/imagenet21k+imagenet2012/\n[`gs://vit_models/augreg`]: https://console.cloud.google.com/storage/browser/vit_models/augreg/\n[`gs://vit_models/sam`]: https://console.cloud.google.com/storage/browser/vit_models/sam/\n[`gs://mixer_models/sam`]: https://console.cloud.google.com/storage/mixer_models/sam/\n[`gs://vit_models/gsam`]: https://console.cloud.google.com/storage/browser/vit_models/gsam/\n[`gs://mixer_models/gsam`]: https://console.cloud.google.com/storage/mixer_models/gsam/\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-research%2Fvision_transformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fgoogle-research%2Fvision_transformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fgoogle-research%2Fvision_transformer/lists"}