Upload 304 files

big_vision_repo/.gitignore

@@ -0,0 +1 @@
+__pycache__

big_vision_repo/CONTRIBUTING.md

@@ -0,0 +1,26 @@
+# How to Contribute
+
+At this time we do not plan to accept non-trivial contributions. The main
+purpose of this codebase is to allow the community to reproduce results from our
+publications.
+
+You are however free to start a fork of the project for your purposes as
+permitted by the license.
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement (CLA). You (or your employer) retain the copyright to your
+contribution; this simply gives us permission to use and redistribute your
+contributions as part of the project. Head over to
+<https://cla.developers.google.com/> to see your current agreements on file or
+to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Community Guidelines
+
+This project follows
+[Google's Open Source Community Guidelines](https://opensource.google/conduct/).

big_vision_repo/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

big_vision_repo/README.md

@@ -0,0 +1,498 @@
+# Big Vision
+
+This codebase is designed for training large-scale vision models using
+[Cloud TPU VMs](https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms)
+or GPU machines. It is based on [Jax](https://github.com/google/jax)/[Flax](https://github.com/google/flax)
+libraries, and uses [tf.data](https://www.tensorflow.org/guide/data) and
+[TensorFlow Datasets](https://www.tensorflow.org/datasets) for scalable and
+reproducible input pipelines.
+
+The open-sourcing of this codebase has two main purposes:
+1. Publishing the code of research projects developed in this codebase (see a
+   list below).
+2. Providing a strong starting point for running large-scale vision experiments
+   on GPU machines and Google Cloud TPUs, which should scale seamlessly and
+   out-of-the box from a single TPU core to a distributed setup with up to 2048
+   TPU cores.
+
+`big_vision` aims to support research projects at Google. We are unlikely to
+work on feature requests or accept external contributions, unless they were
+pre-approved (ask in an issue first). For a well-supported transfer-only
+codebase, see also [vision_transformer](https://github.com/google-research/vision_transformer).
+
+Note that `big_vision` is quite dynamic codebase and, while we intend to keep
+the core code fully-functional at all times, we can not guarantee timely updates
+of the project-specific code that lives in the `.../proj/...` subfolders.
+However, we provide a [table](#project-specific-commits) with last known
+commits where specific projects were known to work.
+
+The following research projects were originally conducted in the `big_vision`
+codebase:
+
+### Architecture research
+
+- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929), by
+  Alexey Dosovitskiy*, Lucas Beyer*, Alexander Kolesnikov*, Dirk Weissenborn*,
+  Xiaohua Zhai*, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer,
+  Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby*
+- [Scaling Vision Transformers](https://arxiv.org/abs/2106.04560), by
+  Xiaohua Zhai*, Alexander Kolesnikov*, Neil Houlsby, and Lucas Beyer*\
+  Resources: [config](big_vision/configs/proj/scaling_laws/train_vit_g.py).
+- [How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers](https://arxiv.org/abs/2106.10270), by
+  Andreas Steiner*, Alexander Kolesnikov*, Xiaohua Zhai*, Ross Wightman,
+  Jakob Uszkoreit, and Lucas Beyer*
+- [MLP-Mixer: An all-MLP Architecture for Vision](https://arxiv.org/abs/2105.01601), by
+  Ilya Tolstikhin*, Neil Houlsby*, Alexander Kolesnikov*, Lucas Beyer*,
+  Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner,
+  Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy\
+  Resources: [config](big_vision/configs/mlp_mixer_i1k.py).
+- [Better plain ViT baselines for ImageNet-1k](https://arxiv.org/abs/2205.01580), by
+  Lucas Beyer, Xiaohua Zhai, Alexander Kolesnikov\
+  Resources: [config](big_vision/configs/vit_s16_i1k.py)
+- [UViM: A Unified Modeling Approach for Vision with Learned Guiding Codes](https://arxiv.org/abs/2205.10337), by
+  Alexander Kolesnikov^*, André Susano Pinto^*, Lucas Beyer*, Xiaohua Zhai*, Jeremiah Harmsen*, Neil Houlsby*\
+  Resources: [readme](big_vision/configs/proj/uvim/README.md), [configs](big_vision/configs/proj/uvim), [colabs](big_vision/configs/proj/uvim).
+- [FlexiViT: One Model for All Patch Sizes](https://arxiv.org/abs/2212.08013), by
+  Lucas Beyer*, Pavel Izmailov*, Alexander Kolesnikov*, Mathilde Caron*, Simon
+  Kornblith*, Xiaohua Zhai*, Matthias Minderer*, Michael Tschannen*, Ibrahim
+  Alabdulmohsin*, Filip Pavetic*\
+  Resources: [readme](big_vision/configs/proj/flexivit/README.md), [configs](big_vision/configs/proj/flexivit).
+- [Dual PatchNorm](https://arxiv.org/abs/2302.01327), by Manoj Kumar, Mostafa Dehghani, Neil Houlsby.
+- [Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design](https://arxiv.org/abs/2305.13035), by
+  Ibrahim Alabdulmohsin*, Xiaohua Zhai*, Alexander Kolesnikov, Lucas Beyer*.
+- (partial) [Scaling Vision Transformers to 22 Billion Parameters](https://arxiv.org/abs/2302.05442), by
+  Mostafa Dehghani*, Josip Djolonga*, Basil Mustafa*, Piotr Padlewski*, Jonathan Heek*, *wow many middle authors*, Neil Houlsby*.
+- (partial) [Finite Scalar Quantization: VQ-VAE Made Simple](https://arxiv.org/abs/2309.15505), by
+  Fabian Mentzer, David Minnen, Eirikur Agustsson, Michael Tschannen.
+- [GIVT: Generative Infinite-Vocabulary Transformers](https://arxiv.org/abs/2312.02116), by
+  Michael Tschannen, Cian Eastwood, Fabian Mentzer.\
+  Resources: [readme](big_vision/configs/proj/givt/README.md), [config](big_vision/configs/proj/givt/givt_imagenet2012.py), [colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/givt/givt_demo_colab.ipynb).
+- [Unified Auto-Encoding with Masked Diffusion](https://arxiv.org/abs/2406.17688), by
+  Philippe Hansen-Estruch, Sriram Vishwanath, Amy Zhang, Manan Tomar.
+
+
+### Multimodal research
+
+- [LiT: Zero-Shot Transfer with Locked-image Text Tuning](https://arxiv.org/abs/2111.07991), by
+  Xiaohua Zhai*, Xiao Wang*, Basil Mustafa*, Andreas Steiner*, Daniel Keysers,
+  Alexander Kolesnikov, and Lucas Beyer*\
+  Resources: [trainer](big_vision/trainers/proj/image_text/contrastive.py), [config](big_vision/configs/proj/image_text/lit_coco.py), [colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb).
+- [Image-and-Language Understanding from Pixels Only](https://arxiv.org/abs/2212.08045), by
+  Michael Tschannen, Basil Mustafa, Neil Houlsby\
+  Resources: [readme](big_vision/configs/proj/clippo/README.md), [config](big_vision/configs/proj/clippo/train_clippo.py), [colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/clippo/clippo_colab.ipynb).
+- [Sigmoid Loss for Language Image Pre-Training](https://arxiv.org/abs/2303.15343), by
+  Xiaohua Zhai*, Basil Mustafa, Alexander Kolesnikov, Lucas Beyer*\
+  Resources: [colab and models](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP_demo.ipynb), code TODO.
+- [A Study of Autoregressive Decoders for Multi-Tasking in Computer Vision](https://arxiv.org/abs/2303.17376), by
+  Lucas Beyer*, Bo Wan*, Gagan Madan*, Filip Pavetic*, Andreas Steiner*, Alexander Kolesnikov, André Susano Pinto, Emanuele Bugliarello, Xiao Wang, Qihang Yu, Liang-Chieh Chen, Xiaohua Zhai*.
+- [Image Captioners Are Scalable Vision Learners Too](https://arxiv.org/abs/2306.07915), by
+  Michael Tschannen*, Manoj Kumar*, Andreas Steiner*, Xiaohua Zhai, Neil Houlsby, Lucas Beyer*.\
+  Resources: [readme](big_vision/configs/proj/cappa/README.md), [config](big_vision/configs/proj/cappa/pretrain.py), [model](big_vision/models/proj/cappa/cappa.py).
+- [Three Towers: Flexible Contrastive Learning with Pretrained Image Models](https://arxiv.org/abs/2305.16999), by Jannik Kossen, Mark Collier, Basil Mustafa, Xiao Wang, Xiaohua Zhai, Lucas Beyer, Andreas Steiner, Jesse Berent, Rodolphe Jenatton, Efi Kokiopoulou.
+- (partial) [PaLI: A Jointly-Scaled Multilingual Language-Image Model](https://arxiv.org/abs/2209.06794), by Xi Chen, Xiao Wang, Soravit Changpinyo, *wow so many middle authors*, Anelia Angelova, Xiaohua Zhai, Neil Houlsby, Radu Soricut.
+- (partial) [PaLI-3 Vision Language Models: Smaller, Faster, Stronger](https://arxiv.org/abs/2310.09199), by Xi Chen, Xiao Wang, Lucas Beyer, Alexander Kolesnikov, Jialin Wu, Paul Voigtlaender, Basil Mustafa, Sebastian Goodman, Ibrahim Alabdulmohsin, Piotr Padlewski, Daniel Salz, Xi Xiong, Daniel Vlasic, Filip Pavetic, Keran Rong, Tianli Yu, Daniel Keysers, Xiaohua Zhai, Radu Soricut.
+- [LocCa](https://arxiv.org/abs/2403.19596), by
+  Bo Wan, Michael Tschannen, Yongqin Xian, Filip Pavetic, Ibrahim Alabdulmohsin, Xiao Wang, André Susano Pinto, Andreas Steiner, Lucas Beyer, Xiaohua Zhai.
+- [PaliGemma](https://arxiv.org/abs/2407.07726), by *wow many authors*.\
+- Resources: [readme](big_vision/configs/proj/paligemma/README.md),
+    [model](big_vision/models/proj/paligemma/paligemma.py),
+    [transfer configs](big_vision/configs/proj/paligemma/transfers),
+    [datasets](big_vision/datasets),
+    [CountBenchQA](big_vision/datasets/countbenchqa/data/countbench_paired_questions.json).
+
+### Training
+
+- [Knowledge distillation: A good teacher is patient and consistent](https://arxiv.org/abs/2106.05237), by
+  Lucas Beyer*, Xiaohua Zhai*, Amélie Royer*, Larisa Markeeva*, Rohan Anil,
+  and Alexander Kolesnikov*\
+  Resources: [README](big_vision/configs/proj/distill/README.md), [trainer](big_vision/trainers/proj/distill/distill.py), [colab](https://colab.research.google.com/drive/1nMykzUzsfQ_uAxfj3k35DYsATnG_knPl?usp=sharing).
+- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://arxiv.org/abs/2010.01412), by
+  Pierre Foret, Ariel Kleiner, Hossein Mobahi, Behnam Neyshabur
+- [Surrogate Gap Minimization Improves Sharpness-Aware Training](https://arxiv.org/abs/2203.08065), by Juntang Zhuang, Boqing Gong, Liangzhe Yuan, Yin Cui, Hartwig Adam, Nicha Dvornek, Sekhar Tatikonda, James Duncan and Ting Liu \
+  Resources: [trainer](big_vision/trainers/proj/gsam/gsam.py), [config](big_vision/configs/proj/gsam/vit_i1k_gsam_no_aug.py) [reproduced results](https://github.com/google-research/big_vision/pull/8#pullrequestreview-1078557411)
+- [Tuning computer vision models with task rewards](https://arxiv.org/abs/2302.08242), by
+  André Susano Pinto*, Alexander Kolesnikov*, Yuge Shi, Lucas Beyer, Xiaohua Zhai.
+- (partial) [VeLO: Training Versatile Learned Optimizers by Scaling Up](https://arxiv.org/abs/2211.09760) by
+  Luke Metz, James Harrison, C. Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, Jascha Sohl-Dickstein.
+
+### Misc
+
+- [Are we done with ImageNet?](https://arxiv.org/abs/2006.07159), by
+  Lucas Beyer*, Olivier J. Hénaff*, Alexander Kolesnikov*, Xiaohua Zhai*, Aäron van den Oord*.
+- [No Filter: Cultural and Socioeconomic Diversity in Contrastive Vision-Language Models](https://arxiv.org/abs/2405.13777), by
+  Angéline Pouget, Lucas Beyer, Emanuele Bugliarello, Xiao Wang, Andreas Peter Steiner, Xiaohua Zhai, Ibrahim Alabdulmohsin.
+
+# Codebase high-level organization and principles in a nutshell
+
+The main entry point is a trainer module, which typically does all the
+boilerplate related to creating a model and an optimizer, loading the data,
+checkpointing and training/evaluating the model inside a loop. We provide the
+canonical trainer `train.py` in the root folder. Normally, individual projects
+within `big_vision` fork and customize this trainer.
+
+All models, evaluators and preprocessing operations live in the corresponding
+subdirectories and can often be reused between different projects. We encourage
+compatible APIs within these directories to facilitate reusability, but it is
+not strictly enforced, as individual projects may need to introduce their custom
+APIs.
+
+We have a powerful configuration system, with the configs living in the
+`configs/` directory. Custom trainers and modules can directly extend/modify
+the configuration options.
+
+Project-specific code resides in the `.../proj/...` namespace. It is not always
+possible to keep project-specific in sync with the core `big_vision` libraries,
+Below we provide the [last known commit](#project-specific-commits)
+for each project where the project code is expected to work.
+
+Training jobs are robust to interruptions and will resume seamlessly from the
+last saved checkpoint (assuming a user provides the correct `--workdir` path).
+
+Each configuration file contains a comment at the top with a `COMMAND` snippet
+to run it, and some hint of expected runtime and results. See below for more
+details, but generally speaking, running on a GPU machine involves calling
+`python -m COMMAND` while running on TPUs, including multi-host, involves
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all
+  --command "bash big_vision/run_tpu.sh COMMAND"
+```
+
+See instructions below for more details on how to run `big_vision` code on a
+GPU machine or Google Cloud TPU.
+
+By default we write checkpoints and logfiles. The logfiles are a list of JSON
+objects, and we provide a short and straightforward [example colab to read
+and display the logs and checkpoints](https://colab.research.google.com/drive/1R_lvV542WUp8Q2y8sbyooZOGCplkn7KI?usp=sharing).
+
+# Current and future contents
+
+The first release contains the core part of pre-training, transferring, and
+evaluating classification models at scale on Cloud TPU VMs.
+
+We have since added the following key features and projects:
+- Contrastive Image-Text model training and evaluation as in LiT and CLIP.
+- Patient and consistent distillation.
+- Scaling ViT.
+- MLP-Mixer.
+- UViM.
+
+Features and projects we plan to release in the near future, in no particular
+order:
+- ImageNet-21k in TFDS.
+- Loading misc public models used in our publications (NFNet, MoCov3, DINO).
+- Memory-efficient Polyak-averaging implementation.
+- Advanced JAX compute and memory profiling. We are using internal tools for
+    this, but may eventually add support for the publicly available ones.
+
+We will continue releasing code of our future publications developed within
+`big_vision` here.
+
+### Non-content
+
+The following exist in the internal variant of this codebase, and there is no
+plan for their release:
+- Regular regression tests for both quality and speed. They rely heavily on
+    internal infrastructure.
+- Advanced logging, monitoring, and plotting of experiments. This also relies
+    heavily on internal infrastructure. However, we are open to ideas on this
+    and may add some in the future, especially if implemented in a
+    self-contained manner.
+- Not yet published, ongoing research projects.
+
+
+# GPU Setup
+
+We first discuss how to setup and run `big_vision` on a (local) GPU machine,
+and then discuss the setup for Cloud TPUs. Note that data preparation step for
+(local) GPU setup can be largely reused for the Cloud TPU setup. While the
+instructions skip this for brevity, we highly recommend using a
+[virtual environment](https://docs.python.org/3/library/venv.html) when
+installing python dependencies.
+
+## Setting up python packages
+
+The first step is to checkout `big_vision` and install relevant python
+dependencies:
+
+```
+git clone https://github.com/google-research/big_vision
+cd big_vision/
+pip3 install --upgrade pip
+pip3 install -r big_vision/requirements.txt
+```
+
+The latest version of `jax` library can be fetched as
+
+```
+pip3 install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+
+You may need a different `jax` package, depending on CUDA and cuDNN libraries
+installed on your machine. Please consult
+[official jax documentation](https://github.com/google/jax#pip-installation-gpu-cuda)
+for more information.
+
+## Preparing tfds data
+
+For unified and reproducible access to standard datasets we opted to use the
+`tensorflow_datasets` (`tfds`) library. It requires each dataset to be
+downloaded, preprocessed and then to be stored on a hard drive (or, if you use
+"Google Cloud", preferably stored in a "GCP bucket".).
+
+Many datasets can be downloaded and preprocessed automatically when used
+for the first time. Nevertheless, we intentionally disable this feature and
+recommend doing dataset preparation step separately, ahead of the first run. It
+will make debugging easier if problems arise and some datasets, like
+`imagenet2012`, require manually downloaded data.
+
+Most of the datasets, e.g. `cifar100`, `oxford_iiit_pet` or `imagenet_v2`
+can be fully automatically downloaded and prepared by running
+
+```
+cd big_vision/
+python3 -m big_vision.tools.download_tfds_datasets cifar100 oxford_iiit_pet imagenet_v2
+```
+
+A full list of datasets is available at [this link](https://www.tensorflow.org/datasets/catalog/overview#all_datasets).
+
+Some datasets, like `imagenet2012` or `imagenet2012_real`, require the data to
+be downloaded manually and placed into `$TFDS_DATA_DIR/downloads/manual/`,
+which defaults to `~/tensorflow_datasets/downloads/manual/`. For example, for
+`imagenet2012` and `imagenet2012_real` one needs to place the official
+`ILSVRC2012_img_train.tar` and `ILSVRC2012_img_val.tar` files in that directory
+and then run
+`python3 -m big_vision.tools.download_tfds_datasets imagenet2012 imagenet2012_real`
+(which may take ~1 hour).
+
+If you use `Google Cloud` and, TPUs in particular, you can then upload
+the preprocessed data (stored in `$TFDS_DATA_DIR`) to
+"Google Cloud Bucket" and use the bucket on any of your (TPU) virtual
+machines to access the data.
+
+## Running on a GPU machine
+
+Finally, after installing all python dependencies and preparing `tfds` data,
+the user can run the job using config of their choice, e.g. to train `ViT-S/16`
+model on ImageNet data, one should run the following command:
+
+```
+python3 -m big_vision.train --config big_vision/configs/vit_s16_i1k.py --workdir workdirs/`date '+%m-%d_%H%M'`
+```
+
+or to train MLP-Mixer-B/16, run (note the `gpu8` config param that reduces the default batch size and epoch count):
+
+```
+python3 -m big_vision.train --config big_vision/configs/mlp_mixer_i1k.py:gpu8 --workdir workdirs/`date '+%m-%d_%H%M'`
+```
+
+# Cloud TPU VM setup
+
+## Create TPU VMs
+
+To create a single machine with 8 TPU cores, follow the following Cloud TPU JAX
+document:
+https://cloud.google.com/tpu/docs/run-calculation-jax
+
+To support large-scale vision research, more cores with multiple hosts are
+recommended. Below we provide instructions on how to do it.
+
+First, create some useful variables, which we be reused:
+
+```
+export NAME=<a name of the TPU deployment, e.g. my-tpu-machine>
+export ZONE=<GCP geographical zone, e.g. europe-west4-a>
+export GS_BUCKET_NAME=<Name of the storage bucket, e.g. my_bucket>
+```
+
+The following command line will create TPU VMs with 32 cores,
+4 hosts.
+
+```
+gcloud compute tpus tpu-vm create $NAME --zone $ZONE --accelerator-type v3-32 --version tpu-ubuntu2204-base
+```
+
+## Install `big_vision` on TPU VMs
+
+Fetch the `big_vision` repository, copy it to all TPU VM hosts, and install
+dependencies.
+
+```
+git clone https://github.com/google-research/big_vision
+gcloud compute tpus tpu-vm scp --recurse big_vision/big_vision $NAME: --zone=$ZONE --worker=all
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "bash big_vision/run_tpu.sh"
+```
+
+## Download and prepare TFDS datasets
+
+We recommend preparing `tfds` data locally as described above and then uploading
+the data to `Google Cloud` bucket. However, if you prefer, the datasets which
+do not require manual downloads can be prepared automatically using a TPU
+machine as described below. Note that TPU machines have only 100 GB of disk
+space, and multihost TPU slices do not allow for external disks to be attached
+in a write mode, so the instructions below may not work for preparing large
+datasets. As yet another alternative, we provide instructions
+[on how to prepare `tfds` data on CPU-only GCP machine](#preparing-tfds-data-on-a-standalone-gcp-cpu-machine).
+
+Specifically, the seven TFDS datasets used during evaluations will be generated
+under `~/tensorflow_datasets` on TPU machine with this command:
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=0 --command "TFDS_DATA_DIR=~/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.tools.download_tfds_datasets cifar10 cifar100 oxford_iiit_pet oxford_flowers102 cars196 dtd uc_merced"
+```
+
+You can then copy the datasets to GS bucket, to make them accessible to all TPU workers.
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=0 --command "rm -r ~/tensorflow_datasets/downloads && gsutil cp -r ~/tensorflow_datasets gs://$GS_BUCKET_NAME"
+```
+
+If you want to integrate other public or custom datasets, i.e. imagenet2012,
+please follow [the official guideline](https://www.tensorflow.org/datasets/catalog/overview).
+
+## Pre-trained models
+
+For the full list of pre-trained models check out the `load` function defined in
+the same module as the model code. And for example config on how to use these
+models, see `configs/transfer.py`.
+
+## Run the transfer script on TPU VMs
+
+The following command line fine-tunes a pre-trained `vit-i21k-augreg-b/32` model
+on `cifar10` dataset.
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/transfer.py:model=vit-i21k-augreg-b/32,dataset=cifar10,crop=resmall_crop --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'` --config.lr=0.03"
+```
+
+## Run the train script on TPU VMs
+
+To train your own big_vision models on a large dataset,
+e.g. `imagenet2012` ([prepare the TFDS dataset](https://www.tensorflow.org/datasets/catalog/imagenet2012)),
+run the following command line.
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/bit_i1k.py  --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'`"
+```
+
+## FSDP training.
+
+`big_vision` supports flexible parameter and model sharding strategies.
+Currently, we support a popular FSDP sharding via a simple config change, see [this config example](big_vision/configs/transfer.py).
+For example, to run FSDP finetuning of a pretrained ViT-L model, run the following command (possible adjusting batch size depending on your hardware):
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.train --config big_vision/configs/transfer.py:model=vit-i21k-augreg-l/16,dataset=oxford_iiit_pet,crop=resmall_crop,fsdp=True,batch_size=256 --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'` --config.lr=0.03"
+```
+
+## Image-text training with SigLIP.
+
+A minimal example that uses public `coco` captions data:
+
+```
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.trainers.proj.image_text.siglip --config big_vision/configs/proj/image_text/siglip_lit_coco.py --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%Y-%m-%d_%H%M'`"
+```
+
+
+
+## Sometimes useful gcloud commands
+
+- Destroy the TPU machines: `gcloud compute tpus tpu-vm delete $NAME --zone $ZONE`
+- Remove all big_vision-related folders on all hosts: `gcloud compute tpus tpu-vm ssh $NAME --zone $ZONE --worker=all --command 'rm -rf ~/big_vision ~/bv_venv'`
+
+## Preparing `tfds` data on a standalone GCP CPU machine.
+
+First create a new machine and a disk (feel free to adjust exact machine type and disk settings/capacity):
+
+```
+export NAME_CPU_HOST=<A name of a CPU-only machine>
+export NAME_DISK=<A name of a disk>
+gcloud compute instances create $NAME_CPU_HOST --machine-type c3-standard-22 --zone $ZONE --image-family ubuntu-2204-lts --image-project ubuntu-os-cloud
+gcloud compute disks create $NAME_DISK --size 1000GB --zone $ZONE --type pd-balanced
+```
+
+Now attach the disk to the newly create machine:
+
+```
+gcloud compute instances attach-disk $NAME_CPU_HOST --disk $NAME_DISK --zone $ZONE
+```
+
+Next, `ssh` to the machine `gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE` and
+[follow instructions to format and mount the disk](https://cloud.google.com/compute/docs/disks/format-mount-disk-linux).
+Let's assume it was mounted to `/mnt/disks/tfds`.
+
+Almost there, now clone and set up `big_vision`:
+
+```
+gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "git clone https://github.com/google-research/big_vision.git && cd big_vision && sh big_vision/run_tpu.sh"
+```
+
+Finally, prepare the dataset (e.g. `coco_captions`) using the utility script and
+copy the result to you google cloud bucket:
+
+```
+gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "cd big_vision && TFDS_DATA_DIR=/mnt/disks/tfds/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.tools.download_tfds_datasets coco_captions"
+gcloud compute ssh $NAME_CPU_HOST --zone=$ZONE --command "rm -rf /mnt/disks/tfds/tensorflow_datasets/downloads && gsutil cp -r /mnt/disks/tfds/tensorflow_datasets gs://$GS_BUCKET_NAME"
+```
+
+
+# ViT baseline
+
+We provide a well-tuned ViT-S/16 baseline in the config file named
+`vit_s16_i1k.py`. It achieves 76.5% accuracy on ImageNet validation split in
+90 epochs of training, being a strong and simple starting point for research
+on the ViT models.
+
+Please see our [arXiv note](https://arxiv.org/abs/2205.01580) for more details
+and if this baseline happens to by useful for your research, consider citing
+
+```
+@article{vit_baseline,
+  url = {https://arxiv.org/abs/2205.01580},
+  author = {Beyer, Lucas and Zhai, Xiaohua and Kolesnikov, Alexander},
+  title = {Better plain ViT baselines for ImageNet-1k},
+  journal={arXiv preprint arXiv:2205.01580},
+  year = {2022},
+}
+```
+
+# Project specific commits
+
+The last known commit where the specific project code is expected to work. The
+core code and configs are expected to work at head.
+
+| Project    | Commit                                                                                        |
+|------------|-----------------------------------------------------------------------------------------------|
+| UViM       | https://github.com/google-research/big_vision/commit/21bd6ebe253f070f584d8b777ad76f4abce51bef |
+| image_text | https://github.com/google-research/big_vision/commit/8921d5141504390a8a4f7b2dacb3b3c042237290 |
+| distill    | https://github.com/google-research/big_vision/commit/2f3f493af048dbfd97555ff6060f31a0e686f17f |
+| GSAM       | WIP                                                                                           |
+| CLIPPO     | https://github.com/google-research/big_vision/commit/fd2d3bd2efc9d89ea959f16cd2f58ae8a495cd44 |
+| CapPa      | https://github.com/google-research/big_vision/commit/7ace659452dee4b68547575352c022a2eef587a5 |
+| GIVT       | https://github.com/google-research/big_vision/commit/0cb70881dd33b3343b769347dc19793c4994b8cb |
+
+# Citing the codebase
+
+If you found this codebase useful for your research, please consider using
+the following BibTEX to cite it:
+
+```
+@misc{big_vision,
+  author = {Beyer, Lucas and Zhai, Xiaohua and Kolesnikov, Alexander},
+  title = {Big Vision},
+  year = {2022},
+  publisher = {GitHub},
+  journal = {GitHub repository},
+  howpublished = {\url{https://github.com/google-research/big_vision}}
+}
+```
+
+# Disclaimer
+
+This is not an official Google Product.
+
+# License
+
+Unless explicitly noted otherwise, everything in the big_vision codebase
+(including models and colabs) is released under the Apache2 license.
+See the LICENSE file for the full license text.

big_vision_repo/big_vision/configs/bit_i1k.py

@@ -0,0 +1,102 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Pre-training BiT on ILSVRC-2012 as in https://arxiv.org/abs/1912.11370
+
+Run training of a BiT-ResNet-50x1 variant, which takes ~32min on v3-128:
+
+big_vision.train \
+    --config big_vision/configs/bit_i1k.py \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+    --config.model.depth 50 --config.model.width 1
+"""
+
+# from big_vision.configs.common_fewshot import get_fewshot_lsr
+import ml_collections as mlc
+
+
+def get_config(runlocal=False):
+  """Config for training on ImageNet-1k."""
+  config = mlc.ConfigDict()
+
+  config.seed = 0
+  config.total_epochs = 90
+  config.num_classes = 1000
+  config.loss = 'softmax_xent'
+
+  config.input = dict()
+  config.input.data = dict(
+      name='imagenet2012',
+      split='train[:99%]',
+  )
+  config.input.batch_size = 4096
+  config.input.cache_raw = True  # Needs up to 120GB of RAM!
+  config.input.shuffle_buffer_size = 250_000  # Per host.
+
+  pp_common = '|onehot(1000, key="{lbl}", key_result="labels")'
+  pp_common += '|value_range(-1, 1)|keep("image", "labels")'
+  config.input.pp = 'decode_jpeg_and_inception_crop(224)|flip_lr' + pp_common.format(lbl='label')
+  pp_eval = 'decode|resize_small(256)|central_crop(224)' + pp_common
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+
+  # Model section
+  config.model_name = 'bit'
+  config.model = dict(
+      depth=50,  # You can also pass e.g. [3, 5, 10, 2]
+      width=1.0,
+  )
+
+  # Optimizer section
+  config.optax_name = 'big_vision.momentum_hp'
+  config.grad_clip_norm = 1.0
+
+  # linear scaling rule. Don't forget to sweep if sweeping batch_size.
+  config.wd = (1e-4 / 256) * config.input.batch_size
+  config.lr = (0.1 / 256) * config.input.batch_size
+  config.schedule = dict(decay_type='cosine', warmup_steps=1000)
+
+  # Eval section
+  def get_eval(split, dataset='imagenet2012'):
+    return dict(
+        type='classification',
+        data=dict(name=dataset, split=split),
+        pp_fn=pp_eval.format(lbl='label'),
+        loss_name=config.loss,
+        log_steps=1000,  # Very fast O(seconds) so it's fine to run it often.
+        cache='final_data',
+    )
+  config.evals = {}
+  config.evals.train = get_eval('train[:2%]')
+  config.evals.minival = get_eval('train[99%:]')
+  config.evals.val = get_eval('validation')
+  config.evals.v2 = get_eval('test', dataset='imagenet_v2')
+  config.evals.real = get_eval('validation', dataset='imagenet2012_real')
+  config.evals.real.pp_fn = pp_eval.format(lbl='real_label')
+
+  # config.evals.fewshot = get_fewshot_lsr(runlocal=runlocal)
+  # config.evals.fewshot.log_steps = 1000
+
+  if runlocal:
+    config.input.batch_size = 32
+    config.input.cache_raw = False
+    config.input.shuffle_buffer_size = 100
+
+    local_eval = config.evals.val
+    config.evals = {'val': local_eval}
+    config.evals.val.cache = 'none'
+
+  return config

big_vision_repo/big_vision/configs/bit_i21k.py

@@ -0,0 +1,85 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""A config for pre-training BiT on ImageNet-21k.
+
+This config relies on the Imagenet-21k tfds dataset, which is not yet
+available publicly in TFDS. We intend to add the dataset to public TFDS soon,
+and this config will then be runnable.
+"""
+
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+import ml_collections as mlc
+
+
+def get_config():
+  """Config for training on imagenet-21k."""
+  config = mlc.ConfigDict()
+
+  config.seed = 0
+  config.total_epochs = 90
+  config.num_classes = 21843
+  config.init_head_bias = -10.0
+  config.loss = 'sigmoid_xent'
+
+  config.input = dict()
+  config.input.data = dict(
+      name='imagenet21k',
+      split='full[51200:]',
+  )
+  config.input.batch_size = 4096
+  config.input.shuffle_buffer_size = 250_000  # Per host, so small-ish is ok.
+
+  pp_common = '|value_range(-1, 1)|onehot({onehot_args})|keep("image", "labels")'
+  pp_common_i21k = pp_common.format(onehot_args=f'{config.num_classes}')
+  pp_common_i1k = pp_common.format(onehot_args='1000, key="label", key_result="labels"')
+  config.input.pp = 'decode_jpeg_and_inception_crop(224)|flip_lr' + pp_common_i21k
+  pp_eval = 'decode|resize_small(256)|central_crop(224)'
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+
+  # Model section
+  config.model_name = 'bit_paper'
+  config.model = dict(depth=50, width=1.0)
+
+  # Optimizer section
+  config.optax_name = 'big_vision.momentum_hp'
+  config.grad_clip_norm = 1.0
+
+  # linear scaling rule. Don't forget to sweep if sweeping batch_size.
+  config.lr = (0.03 / 256) * config.input.batch_size
+  config.wd = (3e-5 / 256) * config.input.batch_size
+  config.schedule = dict(decay_type='cosine', warmup_steps=5000)
+
+  # Evaluations on i21k itself.
+  def eval_i21k(split):
+    return dict(
+        type='classification',
+        data={**config.input.data, 'split': split},
+        pp_fn=pp_eval + pp_common_i21k,
+        loss_name=config.loss,
+        log_steps=1000,  # Very fast O(seconds) so it's fine to run it often.
+    )
+  config.evals = {}
+  config.evals.test = eval_i21k('full[:25_600]')
+  config.evals.val = eval_i21k('full[25_600:51_200]')
+  config.evals.train = eval_i21k('full[51_200:76_800]')
+
+  # Few-shot evaluators
+  config.evals.fewshot = get_fewshot_lsr()
+  config.evals.fewshot.log_steps = 25_000
+
+  return config

big_vision_repo/big_vision/configs/common.py

@@ -0,0 +1,188 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A few things commonly used across A LOT of config files."""
+
+import string
+
+import ml_collections as mlc
+
+
+def input_for_quicktest(config_input, quicktest):
+  if quicktest:
+    config_input.batch_size = 8
+    config_input.shuffle_buffer_size = 10
+    config_input.cache_raw = False
+
+
+def parse_arg(arg, lazy=False, **spec):
+  """Makes ConfigDict's get_config single-string argument more usable.
+
+  Example use in the config file:
+
+    import big_vision.configs.common as bvcc
+    def get_config(arg):
+      arg = bvcc.parse_arg(arg,
+          res=(224, int),
+          runlocal=False,
+          schedule='short',
+      )
+
+      # ...
+
+      config.shuffle_buffer = 250_000 if not arg.runlocal else 50
+
+  Ways that values can be passed when launching:
+
+    --config amazing.py:runlocal,schedule=long,res=128
+    --config amazing.py:res=128
+    --config amazing.py:runlocal  # A boolean needs no value for "true".
+    --config amazing.py:runlocal=False  # Explicit false boolean.
+    --config amazing.py:128  # The first spec entry may be passed unnamed alone.
+
+  Uses strict bool conversion (converting 'True', 'true' to True, and 'False',
+    'false', '' to False).
+
+  Args:
+    arg: the string argument that's passed to get_config.
+    lazy: allow lazy parsing of arguments, which are not in spec. For these,
+      the type is auto-extracted in dependence of most complex possible type.
+    **spec: the name and default values of the expected options.
+      If the value is a tuple, the value's first element is the default value,
+      and the second element is a function called to convert the string.
+      Otherwise the type is automatically extracted from the default value.
+
+  Returns:
+    ConfigDict object with extracted type-converted values.
+  """
+  # Normalize arg and spec layout.
+  arg = arg or ''  # Normalize None to empty string
+  spec = {k: get_type_with_default(v) for k, v in spec.items()}
+
+  result = mlc.ConfigDict(type_safe=False)  # For convenient dot-access only.
+
+  # Expand convenience-cases for a single parameter without = sign.
+  if arg and ',' not in arg and '=' not in arg:
+    # (think :runlocal) If it's the name of sth in the spec (or there is no
+    # spec), it's that in bool.
+    if arg in spec or not spec:
+      arg = f'{arg}=True'
+    # Otherwise, it is the value for the first entry in the spec.
+    else:
+      arg = f'{list(spec.keys())[0]}={arg}'
+      # Yes, we rely on Py3.7 insertion order!
+
+  # Now, expand the `arg` string into a dict of keys and values:
+  raw_kv = {raw_arg.split('=')[0]:
+                raw_arg.split('=', 1)[-1] if '=' in raw_arg else 'True'
+            for raw_arg in arg.split(',') if raw_arg}
+
+  # And go through the spec, using provided or default value for each:
+  for name, (default, type_fn) in spec.items():
+    val = raw_kv.pop(name, None)
+    result[name] = type_fn(val) if val is not None else default
+
+  if raw_kv:
+    if lazy:  # Process args which are not in spec.
+      for k, v in raw_kv.items():
+        result[k] = autotype(v)
+    else:
+      raise ValueError(f'Unhandled config args remain: {raw_kv}')
+
+  return result
+
+
+def get_type_with_default(v):
+  """Returns (v, string_to_v_type) with lenient bool parsing."""
+  # For bool, do safe string conversion.
+  if isinstance(v, bool):
+    def strict_bool(x):
+      assert x.lower() in {'true', 'false', ''}
+      return x.lower() == 'true'
+    return (v, strict_bool)
+  # If already a (default, type) tuple, use that.
+  if isinstance(v, (tuple, list)):
+    assert len(v) == 2 and isinstance(v[1], type), (
+        'List or tuple types are currently not supported because we use `,` as'
+        ' dumb delimiter. Contributions (probably using ast) welcome. You can'
+        ' unblock by using a string with eval(s.replace(";", ",")) or similar')
+    return (v[0], v[1])
+  # Otherwise, derive the type from the default value.
+  return (v, type(v))
+
+
+def autotype(x):
+  """Auto-converts string to bool/int/float if possible."""
+  assert isinstance(x, str)
+  if x.lower() in {'true', 'false'}:
+    return x.lower() == 'true'  # Returns as bool.
+  try:
+    return int(x)  # Returns as int.
+  except ValueError:
+    try:
+      return float(x)  # Returns as float.
+    except ValueError:
+      return x  # Returns as str.
+
+
+def pack_arg(**kw):
+  """Packs key-word args as a string to be parsed by `parse_arg()`."""
+  for v in kw.values():
+    assert ',' not in f'{v}', f"Can't use `,` in config_arg value: {v}"
+  return ','.join([f'{k}={v}' for k, v in kw.items()])
+
+
+def arg(**kw):
+  """Use like `add(**bvcc.arg(res=256, foo=bar), lr=0.1)` to pass config_arg."""
+  return {'config_arg': pack_arg(**kw), **kw}
+
+
+def _get_field_ref(config_dict, field_name):
+  path = field_name.split('.')
+  for field in path[:-1]:
+    config_dict = getattr(config_dict, field)
+  return config_dict.get_ref(path[-1])
+
+
+def format_str(format_string, config):
+  """Format string with reference fields from config.
+
+  This makes it easy to build preprocess strings that contain references to
+  fields tha are edited after. E.g.:
+
+  ```
+  config = mlc.ConficDict()
+  config.res = (256, 256)
+  config.pp = bvcc.format_str('resize({res})', config)
+  ...
+  # if config.res is modified (e.g. via sweeps) it will propagate to pp field:
+  config.res = (512, 512)
+  assert config.pp == 'resize((512, 512))'
+  ```
+
+  Args:
+    format_string: string to format with references.
+    config: ConfigDict to get references to format the string.
+
+  Returns:
+    A reference field which renders a string using references to config fields.
+  """
+  output = ''
+  parts = string.Formatter().parse(format_string)
+  for (literal_text, field_name, format_spec, conversion) in parts:
+    assert not format_spec and not conversion
+    output += literal_text
+    if field_name:
+      output += _get_field_ref(config, field_name).to_str()
+  return output

big_vision_repo/big_vision/configs/common_fewshot.py

@@ -0,0 +1,56 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Most common few-shot eval configuration."""
+
+import ml_collections as mlc
+
+
+def get_fewshot_lsr(target_resolution=224, resize_resolution=256,
+                    runlocal=False, **kw):
+  """Returns a standard-ish fewshot eval configuration."""
+  kw.setdefault('representation_layer', 'pre_logits')
+  kw.setdefault('shots', (1, 5, 10, 25))
+  kw.setdefault('l2_reg', 2.0 ** 10)
+  kw.setdefault('num_seeds', 3)
+  kw.setdefault('prefix', '')  # No prefix as we already use a/ z/ and zz/
+
+  # Backward-compatible default:
+  if not any(f'log_{x}' in kw for x in ['steps', 'percent', 'examples', 'epochs']):  # pylint: disable=line-too-long
+    kw['log_steps'] = 25_000
+
+  config = mlc.ConfigDict(kw)
+  config.type = 'fewshot_lsr'
+  config.datasets = {
+      'caltech': ('caltech101', 'train', 'test'),  # copybara:srtip
+      'cars': ('cars196:2.1.0', 'train', 'test'),
+      'cifar100': ('cifar100', 'train', 'test'),
+      'dtd': ('dtd', 'train', 'test'),
+      # The first 65000 ImageNet samples have at least 30 shots per any class.
+      # Commented out by default because needs manual download.
+      # 'imagenet': ('imagenet2012', 'train[:65000]', 'validation'),
+      'pets': ('oxford_iiit_pet', 'train', 'test'),
+      'uc_merced': ('uc_merced', 'train[:1000]', 'train[1000:]'),
+  } if not runlocal else {
+      'pets': ('oxford_iiit_pet', 'train', 'test'),
+  }
+  config.pp_train = (f'decode|resize({resize_resolution})|'
+                     f'central_crop({target_resolution})|'
+                     f'value_range(-1,1)|keep("image", "label")')
+  config.pp_eval = (f'decode|resize({resize_resolution})|'
+                    f'central_crop({target_resolution})|'
+                    f'value_range(-1,1)|keep("image", "label")')
+  config.display_first = [('imagenet', 10)] if not runlocal else [('pets', 10)]
+
+  return config

big_vision_repo/big_vision/configs/load_and_eval.py

@@ -0,0 +1,143 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pytype: disable=not-writable,attribute-error
+# pylint: disable=line-too-long,missing-function-docstring
+r"""A config to load and eval key model using the core train.py.
+
+The runtime varies widely depending on the model, but each one should reproduce
+the corresponding paper's numbers.
+This configuration makes use of the "arg" to get_config to select which model
+to run, so a few examples are given below:
+
+Run and evaluate a BiT-M ResNet-50x1 model that was transferred to i1k:
+
+big_vision.train \
+    --config big_vision/configs/load_and_eval.py:name=bit_paper,batch_size=8 \
+    --config.model_init M-imagenet2012 --config.model.width 1 --config.model.depth 50
+
+Run and evaluate the recommended ViT-B/32 from "how to train your vit" paper:
+
+big_vision.train \
+    --config big_vision/configs/load_and_eval.py:name=vit_i21k,batch_size=8 \
+    --config.model.variant B/32 --config.model_init howto-i21k-B/32
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+
+
+def eval_only(config, batch_size, spec_for_init):
+  """Set a few configs that turn trainer into (almost) eval-only."""
+  config.total_steps = 0
+  config.input = {}
+  config.input.batch_size = batch_size
+  config.input.data = dict(name='bv:dummy', spec=spec_for_init)
+  config.optax_name = 'identity'
+  config.lr = 0.0
+
+  config.mesh = [('data', -1)]
+  config.sharding_strategy = [('params/.*', 'fsdp(axis="data")')]
+  config.sharding_rules = [('act_batch', ('data',))]
+
+  return config
+
+
+def get_config(arg=''):
+  config = bvcc.parse_arg(arg, name='bit_paper', batch_size=4)
+
+  # Make the config eval-only by setting some dummies.
+  eval_only(config, config.batch_size, spec_for_init=dict(
+      image=dict(shape=(224, 224, 3), dtype='float32'),
+  ))
+
+  config.evals = dict(fewshot=get_fewshot_lsr())
+
+  # Just calls the function with the name given as `config`.
+  # Could also be a giant if-block if you're into that kind of thing.
+  globals()[config.name](config)
+  return config
+
+
+def bit_paper(config):
+  config.num_classes = 1000
+
+  config.model_name = 'bit_paper'
+  config.model_init = 'M-imagenet2012'  # M = i21k, -imagenet2012 = fine-tuned
+  config.model = dict(width=1, depth=50)
+
+  def get_eval(split, lbl, dataset='imagenet2012_real'):
+    return dict(
+        type='classification',
+        data=dict(name=dataset, split=split),
+        loss_name='softmax_xent',
+        cache='none',  # Only run once, on low-mem machine.
+        pp_fn=(
+            'decode|resize(384)|value_range(-1, 1)'
+            f'|onehot(1000, key="{lbl}", key_result="labels")'
+            '|keep("image", "labels")'
+        ),
+    )
+  config.evals.test = get_eval('validation', 'original_label')
+  config.evals.real = get_eval('validation', 'real_label')
+  config.evals.v2 = get_eval('test', 'label', 'imagenet_v2')
+
+
+def vit_i1k(config):
+  config.num_classes = 1000
+
+  config.model_name = 'vit'
+  config.model_init = ''  # Will be set in sweep.
+  config.model = dict(variant='S/16', pool_type='gap', posemb='sincos2d',
+                      rep_size=True)
+
+  config.evals.val = dict(
+      type='classification',
+      data=dict(name='imagenet2012', split='validation'),
+      pp_fn='decode|resize_small(256)|central_crop(224)|value_range(-1, 1)|onehot(1000, key="label", key_result="labels")|keep("image", "labels")',
+      loss_name='softmax_xent',
+      cache='none',  # Only run once, on low-mem machine.
+  )
+
+
+def mlp_mixer_i1k(config):
+  config.num_classes = 1000
+
+  config.model_name = 'mlp_mixer'
+  config.model_init = ''  # Will be set in sweep.
+  config.model = dict(variant='L/16')
+
+  config.evals.val = dict(
+      type='classification',
+      data=dict(name='imagenet2012', split='validation'),
+      pp_fn='decode|resize_small(256)|central_crop(224)|value_range(-1, 1)|onehot(1000, key="label", key_result="labels")|keep("image", "labels")',
+      loss_name='softmax_xent',
+      cache='none',  # Only run once, on low-mem machine.
+  )
+
+
+def vit_i21k(config):
+  config.num_classes = 21843
+
+  config.model_name = 'vit'
+  config.model_init = ''  # Will be set in sweep.
+  config.model = dict(variant='B/32', pool_type='tok')
+
+  config.evals.val = dict(
+      type='classification',
+      data=dict(name='imagenet21k', split='full[:51200]'),
+      pp_fn='decode|resize_small(256)|central_crop(224)|value_range(-1, 1)|onehot(21843)|keep("image", "labels")',
+      loss_name='sigmoid_xent',
+      cache='none',  # Only run once, on low-mem machine.
+  )

big_vision_repo/big_vision/configs/mlp_mixer_i1k.py

@@ -0,0 +1,120 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""A config for training MLP-Mixer-B/16 model on ILSVRC-2012 ("ImageNet-1k").
+
+Achieves 76.3% top-1 accuracy on the test split in 2h11m on TPU v3-128
+with 300 epochs. A shorter 60 epochs run is expected to get to 70.5% in 27m.
+
+big_vision.train \
+    --config big_vision/configs/mlp_mixer_i1k.py \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+"""
+
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+import ml_collections as mlc
+
+
+def get_config(mode=None):
+  """Config for training Mixer on i1k."""
+  config = mlc.ConfigDict()
+
+  config.seed = 0
+  config.total_epochs = 300
+  config.num_classes = 1000
+  config.loss = 'sigmoid_xent'
+  config.init_head_bias = -6.9
+
+  config.input = dict()
+  config.input.data = dict(
+      name='imagenet2012',
+      split='train[:99%]',
+  )
+  config.input.batch_size = 4096
+  config.input.cache_raw = True  # Needs up to 120GB of RAM!
+  config.input.shuffle_buffer_size = 250_000
+
+  config.input.pp = (
+      'decode_jpeg_and_inception_crop(224)'
+      '|flip_lr'
+      '|randaug(2,15)'
+      '|value_range(-1, 1)'
+      '|onehot(1000, key="label", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+  pp_eval = (
+      'decode'
+      '|resize_small(256)|central_crop(224)'
+      '|value_range(-1, 1)'
+      '|onehot(1000, key="{lbl}", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+
+  # To continue using the near-defunct randaug op.
+  config.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'archive.randaug']
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+
+  config.prefetch_to_device = 2
+
+  # Model section
+  config.model_name = 'mlp_mixer'
+  config.model = dict()
+  config.model.variant = 'B/16'
+  config.model.stoch_depth = 0.1
+
+  config.mixup = dict(fold_in=None, p=0.5)
+
+  # Optimizer section
+  config.optax_name = 'scale_by_adam'
+  config.grad_clip_norm = 1.
+
+  config.lr = 0.001
+  config.wd = 1e-4
+  config.schedule = dict(
+      decay_type='linear',
+      warmup_steps=10_000,
+      linear_end=1e-5,
+  )
+
+  # Eval section
+  def get_eval(split, dataset='imagenet2012'):
+    return dict(
+        type='classification',
+        data=dict(name=dataset, split=split),
+        pp_fn=pp_eval.format(lbl='label'),
+        loss_name=config.loss,
+        log_steps=2500,  # Very fast O(seconds) so it's fine to run it often.
+        cache_final=mode != 'gpu8',
+    )
+  config.evals = {}
+  config.evals.train = get_eval('train[:2%]')
+  config.evals.minival = get_eval('train[99%:]')
+  config.evals.val = get_eval('validation')
+  config.evals.v2 = get_eval('test', dataset='imagenet_v2')
+  config.evals.real = get_eval('validation', dataset='imagenet2012_real')
+  config.evals.real.pp_fn = pp_eval.format(lbl='real_label')
+
+  config.fewshot = get_fewshot_lsr()
+
+  if mode == 'gpu8':
+    config.total_epochs = 60
+    config.input.batch_size = 512
+    config.input.cache_raw = False
+  if mode == 'regression_test':
+    config.total_epochs = 60
+
+  return config

big_vision_repo/big_vision/configs/proj/cappa/README.md

@@ -0,0 +1,37 @@
+# Image Captioners Are Scalable Vision Learners Too
+
+*by Michael Tschannen, Manoj Kumar, Andreas Steiner, Xiaohua Zhai, Neil Houlsby, Lucas Beyer* [[arxiv]](https://arxiv.org/abs/2306.07915)
+
+![CapPa Architecture](./cappa_architecture.png)
+
+This directory contains a config for training a CapPa model from scratch.
+Note that most models in the paper were trained on a proprietary dataset
+(WebLI), but similar results can be obtained by training on [LAION](https://laion.ai/).
+
+By default, this config trains on COCO captions as this data set is readily
+available in [TFDS](https://www.tensorflow.org/datasets) without manual steps.
+This is not meant to produce a meaningful model, but
+provides a way for the user to run the config out of the box. Please update the
+config with with a TFDS-wrapped variant of your favorite image/text data set to
+train capable models.
+
+After setting up `big_vision` as described in the [main README](https://github.com/google-research/big_vision#cloud-tpu-vm-setup), training can be launched as follows
+
+```
+python -m big_vision.trainers.proj.cappa.generative \
+  --config big_vision/configs/proj/cappa/pretrain.py \
+  --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`
+```
+
+To run the Cap baseline (autoregressive captioning without parallel prediction),
+set `config.model.masked_pred_prob = 0.0`.
+
+### Citation
+```
+@inproceedings{tschannen2023image,
+  title={Image Captioners Are Scalable Vision Learners Too},
+  author={Tschannen, Michael and Kumar, Manoj and Steiner, Andreas and Zhai, Xiaohua and Houlsby, Neil and Beyer, Lucas},
+  booktitle={Neural Information Processing Systems (NeurIPS)},
+  year={2023}
+}
+```

big_vision_repo/big_vision/configs/proj/cappa/pretrain.py

@@ -0,0 +1,140 @@
+# Copyright 2023 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Trains a CapPa model (https://arxiv.org/abs/2306.07915) on coco_captions.
+
+This config is for reference, we never ran a full training on a large
+image/text data set on public infrastructure.
+
+big_vision.trainers.proj.cappa.generative \
+  --config big_vision/configs/proj/cappa/pretrain.py \
+  --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'`
+"""
+
+
+from big_vision.configs import common_fewshot
+import big_vision.configs.common as bvcc
+import ml_collections
+
+
+def get_config(arg=None):
+  """Returns the base config."""
+  config = bvcc.parse_arg(arg,
+                          runlocal=False,
+                          total_steps=366_500,
+                          batch_size=8*1024,
+                          warmup_steps=10_000,
+                          )
+
+  config.evals = {}
+  config.input = {}
+  config.input.batch_size = config.batch_size if not config.runlocal else 8
+  shuffle_buffer_size = 50_000 if not config.runlocal else 50
+
+  res = 224
+  patch_size = 16
+  max_text_tokens = 64
+
+  pp_image = (f'resize({res})|value_range(-1,1)')
+
+  def tokenizer(inkey, outkey):
+    return (f'tokenize(max_len={max_text_tokens}, model="c4_en", '
+            f'eos="sticky", inkey="{inkey}", outkey="{outkey}")')
+
+  pp_coco = (f'decode|{pp_image}|'
+             'coco_captions("captions")|choice(inkey="captions", outkey="text")|'
+             f'{tokenizer("text", "labels")}|keep("image", "labels")')
+  config.input.pp = pp_coco
+
+  # NOTE: "coco_captions" is way too small a dataset to train on. It's simply
+  # used here to serve as a smoke test that the implementation works correctly.
+  config.input.data = dict(name='coco_captions', split='train')  # num_examples=82_783
+  config.input.shuffle_buffer_size = shuffle_buffer_size
+
+  config.evals.val_coco = {
+      'type': 'proj.cappa.perplexity',
+      'pred': 'perplexity',
+      'log_steps': 1000,
+      'data': dict(name='coco_captions', split='val'),  # num_examples=5_000
+      'pp_fn': pp_coco,
+  }
+
+  # Few-shot  metrics
+  config.evals.fewshot = common_fewshot.get_fewshot_lsr(
+      target_resolution=res, resize_resolution=int(256 / 224 * res))
+  config.evals.fewshot.type = 'fewshot_lsr'
+  config.evals.fewshot.log_steps = 5_000 if not config.runlocal else 5
+  config.evals.fewshot.representation_layer = 'pre_logits'
+  config.evals.fewshot.pred = 'enc_rep'
+  config.evals.fewshot.pp_eval = config.evals.fewshot.pp_train
+
+  # NOTE: Scoring of the entire imagenet validation set is rather slow:
+  # ~100 secs / 1k classes / host.
+  config.evals['imagenet/scoring'] = dict(
+      type='proj.cappa.scoring_classifier',
+      pred='score',
+      log_percent=0.1,
+      data=dict(name='imagenet2012', split='validation'),
+      pp_fn=f'decode|{pp_image}|keep("image", "label")',
+      pp_txt=tokenizer('label', 'labels'),
+  )
+
+  for e in config.evals.values():
+    e.skip_first = True
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None  # 10_000
+
+  # Model section
+  config.model_name = 'proj.cappa.cappa'
+  config.model = ml_collections.ConfigDict()
+  config.model.num_layers = 12
+  config.model.num_heads = 12
+  config.model.mlp_dim = 3072
+  config.model.emb_dim = 768
+  config.model.vocab_size = 32_000
+  config.model.patches = (patch_size, patch_size)
+  config.model.seq_len = max_text_tokens
+  config.model.posemb_type = 'learn'
+
+  # Decoder
+  config.model.decoder_num_layers = 6
+  # 0 values here mean to use the same value as for the encoder
+  config.model.decoder_num_heads = 0
+  config.model.decoder_mlp_dim = 0
+  config.model.decoder_emb_dim = 0
+  config.model.dec_dropout_rate = 0.0
+  config.model.masked_pred_prob = 0.75
+  config.model.masking_ratio = 1.0
+  config.model.decoder_bias = False
+
+  config.optax_name = 'big_vision.scale_by_adafactor'
+  config.optax = dict(beta2_cap=0.999)
+  config.grad_clip_norm = 1.0
+  config.label_smoothing = 0.0
+
+  schedule = dict(decay_type='cosine',
+                  warmup_steps=config.warmup_steps
+                  if not config.runlocal else 5)
+
+  # Standard schedule
+  config.lr = 0.001
+  config.wd = 0.0001
+  config.schedule = schedule
+
+  config.seed = 0
+
+  return config

big_vision_repo/big_vision/configs/proj/clippo/README.md

@@ -0,0 +1,85 @@
+## Image-and-Language Understanding from Pixels Only
+
+*by Michael Tschannen, Basil Mustafa, Neil Houlsby* [[arxiv]](https://arxiv.org/abs/2212.08045) [[colab]](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/clippo/clippo_colab.ipynb)
+
+We provide pretrained CLIP with Pixels Only (CLIPPO) models and code to train such models on image/alt-text data sets.
+
+### Pretrained models
+
+Six ViT-B/16 models trained on a mix of [`YFCC-100M`](https://arxiv.org/abs/1503.01817) and [`C4`](https://arxiv.org/abs/1910.10683) (some initialized with an [ImageNet21k-pretrained checkpoint](https://github.com/google-research/vision_transformer#vision-transformer)\) are available.
+These models were trained using the schedules and hyperparameters described in the paper. We use the full `YFCC-100M` data set, sampling one of the available `title/description/tag` annotations at random for each each example. We drop non-descriptive annotations (e.g. descriptions consisting of digits only) following the filtering procedure outlined in the [LiT paper](https://arxiv.org/abs/2303.04671), Appendix E. The preprocessing for the `C4` data is as described in the paper.
+
+The tables below show details about the checkpoints and their performance on Vision & Language benchmarks, and [`GLUE`](https://arxiv.org/abs/1804.07461). We also provide a [colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/clippo/clippo_colab.ipynb) to load the models, compute embeddings, and perform zero-shot classification.
+
+##### Checkpoint details
+
+| model            | training dataset   | #param.   | steps   | checkpoint |
+|:-----------------|:-------------------|:----------|:--------|:-----------|
+| CLIPPO           | YFCC-100M          | 93M       | 250k    | `gs://big_vision/clippo/clippo_b16_yfcc100m.npz` |
+| CLIPPO I21k init | YFCC-100M          | 93M       | 250k    | `gs://big_vision/clippo/clippo_b16_yfcc100m_i21k_init.npz` |
+| CLIPPO I21k init | YFCC-100M + 25%C4  | 93M       | 333k    | `gs://big_vision/clippo/clippo_b16_yfcc100m_i21k_init_25c4.npz` |
+| CLIPPO I21k init | YFCC-100M + 50%C4  | 93M       | 500k    | `gs://big_vision/clippo/clippo_b16_yfcc100m_i21k_init_50c4.npz` |
+| CLIPPO I21k init | YFCC-100M + 75%C4  | 93M       | 500k    | `gs://big_vision/clippo/clippo_b16_yfcc100m_i21k_init_75c4.npz` |
+| CLIPPO           | C4          | 93M       | 250k    | `gs://big_vision/clippo/clippo_b16_100c4.npz` |
+
+##### Vision \& Language results
+
+| model            | training dataset   | ImageNet 10-shot | ImageNet 0-shot | MS-COCO I→T | MS-COCO T→I |
+|:-----------------|:-------------------|-----------:|----------:|--------:|--------:|
+| CLIPPO           | YFCC-100M          |       38.2 |      43.4 |    34.7 |    19.7 |
+| CLIPPO I21k init | YFCC-100M          |       44.7 |      47.4 |    36.1 |    21.3 |
+| CLIPPO I21k init | YFCC-100M + 25%C4  |       43.8 |      44.8 |    33.3 |    19.4 |
+| CLIPPO I21k init | YFCC-100M + 50%C4  |       41.2 |      42.0 |    31.4 |    17.8 |
+| CLIPPO I21k init | YFCC-100M + 75%C4  |       34.5 |      33.4 |    26.6 |    14.6 |
+
+##### GLUE results
+
+| model            | training dataset   | MNLI-M/MM   |   QQP |   QNLI |   SST-2 |   COLA |   STS-B |   MRPC |   RTE |   avg |
+|:-----------------|:-------------------|:------------|------:|-------:|--------:|-------:|--------:|-------:|------:|------:|
+| CLIPPO           | YFCC-100M          | 71.3 / 71.5 |  79.1 |   67.9 |    85.7 |    0.0 |    14.0 |   83.4 |  54.9 |  58.6 |
+| CLIPPO I21k init | YFCC-100M          | 70.0 / 70.1 |  83.7 |   81.6 |    86.1 |    0.0 |    18.5 |   83.0 |  53.1 |  60.7 |
+| CLIPPO I21k init | YFCC-100M + 25%C4  | 75.7 / 75.1 |  85.2 |   83.5 |    89.6 |    0.0 |    82.3 |   82.7 |  52.7 |  69.7 |
+| CLIPPO I21k init | YFCC-100M + 50%C4  | 77.4 / 77.4 |  86.0 |   83.9 |    91.7 |   34.5 |    84.5 |   85.1 |  56.3 |  75.2 |
+| CLIPPO I21k init | YFCC-100M + 75%C4  | 79.8 / 79.1 |  86.5 |   84.3 |    92.0 |   44.5 |    85.3 |   88.2 |  58.5 |  77.6 |
+| CLIPPO           | C4                 | 79.9 / 80.2 |  86.7 |   85.2 |    93.3 |   50.9 |    84.7 |   86.3 |  58.5 |  78.4 |
+
+### Training your own models
+
+To train your own CLIPPO model, please follow the setup instructions in the [`big_vision` main README](https://github.com/google-research/big_vision#cloud-tpu-vm-setup). In the following, we provide the CLIPPO-specific commands required in addition to the setup, assume you are using the Google Cloud TPU setup (potentially with adapted TPU configuration, see table below). If you are using GPUs, please set up your machine directly and only execute the `--command` portions of the commands below from the `big_vision` repository root.
+
+The text rendering preproprocessing function requires manual download of the Unifont .hex files from [Unifoundry](https://unifoundry.com/unifont/) (please follow link for license):
+
+```bash
+gcloud alpha compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all \
+--command "bash big_vision/pp/proj/clippo/download_unifont.sh"
+```
+
+Launch the training by running
+
+```bash
+gcloud alpha compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all \
+--command "TFDS_DATA_DIR=gs://$GS_BUCKET_NAME/tensorflow_datasets bash big_vision/run_tpu.sh big_vision.trainers.proj.image_text.contrastive --config big_vision/configs/proj/clippo/train_clippo.py --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'`"
+```
+
+*Important note:* The input pipeline relies on [TensorFlow Datasets (TFDS)](https://www.tensorflow.org/datasets) which does not provide automatic integration with large image/alt-text datasets out of the box. The above config therefore trains by default on MS-COCO Captions which can be automatically downloaded via TFDS, and additionally initializes the CLIPPO ViT backbone with weights pretrained on ImageNet21k. This setup is not meant to produce good accuracy, but to provide the user with a way to sanity-check their setup. If you want to train on a large data set such as [`LAION-400M`](https://arxiv.org/abs/2111.02114) or [`YFCC-100M`](https://arxiv.org/abs/1503.01817), please follow [these instructions](https://www.tensorflow.org/datasets/add_dataset) to wrap your data set using TFDS, and update the dataset in the config accordingly. Also note that the ImageNet1k evaluations require manual download of the data, see [these instructions](https://github.com/google-research/big_vision#preparing-tfds-data). To train with your own data set and with ImageNet1k-based evaluations, use `--config big_vision/configs/proj/clippo/train_clippo.py:test_with_coco=False,i1k_eval=True` in the command above.
+
+##### Expected results
+
+| train dataset | batch size | #steps | TPU chips | ImageNet 0-shot | MS-COCO I→T | MS-COCO T→I | Config `arg` |
+| :---  | ---:         | ---: | ---: | :---:           | :---:       | :---:       | :---         |
+| *MS-COCO (sanity check)* | 4000 | 400 | 32 v3 | 4.2 | 12.6 | 8.6 | `i1k_eval=True` |
+| LAION-400M | 8192 | 100k |128 v2 | 51.5 | 44.8 | 29.3 | `test_with_coco=False,i1k_eval=True` |
+| LAION-400M | 10240\* | 100k | 128 v3 | 53.6 | 46.7 | 30.3 | `test_with_coco=False,i1k_eval=True` |
+
+\* The experiments in the paper use a batch size of 10240 which requires a memory-optimized ViT implementation to run on 128 TPU v2 chips or 128 TPU v3 chips (in which case the TPU memory capacity allows to increase the batch size beyond 10240).
+
+### Citation
+
+```
+@inproceedings{tschannen2023image,
+  title={Image-and-Language Understanding from Pixels Only},
+  author={Tschannen, Michael and Mustafa, Basil and Houlsby, Neil},
+  booktitle={Proc. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
+  year={2023}
+}
+```

big_vision_repo/big_vision/configs/proj/clippo/train_clippo.py

@@ -0,0 +1,199 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Trains CLIP with Pixels Only (CLIPPO), https://arxiv.org/abs/2212.08045
+
+IMPORTANT NOTE: This config uses coco_captions by default for demonstration
+purposes since the TFDS catalog does not provide any large image/alt-text data
+set; the training will not produce a model with useful accuracy. Please
+replace the data set below (marked by a comment) with an appropriate image/
+alt-text data set wrapped in TFDS (for example LAION-400M) and run the config
+with the suffix `:test_with_coco=False` to train on your data set. Refer to
+the following guide to build a TFDS wrapper for your favorite image/alt-text
+data set:
+https://www.tensorflow.org/datasets/add_dataset
+
+Also note that evaluation on ImageNet requires manual TFDS setup, see
+https://github.com/google-research/big_vision#preparing-tfds-data
+
+
+Example training:
+
+big_vision.trainers.proj.image_text.contrastive \
+    --config big_vision/configs/proj/clippo/train_clippo.py \
+    --workdir gs://[your_bucket]/big_vision/`date '+%Y-%m-%d_%H%M'`
+
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+from big_vision.configs.proj.image_text import common
+from ml_collections import ConfigDict
+
+
+def get_config(arg=None):
+  """The base configuration."""
+  arg = bvcc.parse_arg(
+      arg, res=224, runlocal=False, variant='B/16',
+      test_with_coco=True, i1k_eval=False)
+  config = ConfigDict()
+
+  config.input = {}
+  if arg.test_with_coco:
+    # Use COCO Captions for sanity-checking
+    config.input.data = dict(name='coco_captions', split='train')
+    val_data = dict(config.input.data)
+    val_data['split'] = 'val'
+    config.input.batch_size = 4000 if not arg.runlocal else 32
+    config.input.shuffle_buffer_size = 50_000  if not arg.runlocal else 50
+    config.total_steps = 400 if not arg.runlocal else 10
+  else:
+    # Please add your favorite image/alt-text dataset here
+    config.input.data = None
+    val_data = None
+    assert config.input.data is not None and val_data is not None, (
+        config.input.data, val_data)
+
+    # The value in the paper is 10 * 1024, which requires 128 TPUv3 cores or a
+    # memory optimized ViT implementation when running on 128 TPUv2 cores.
+    config.input.batch_size = 8 * 1024 if not arg.runlocal else 32
+    config.input.shuffle_buffer_size = 250_000  if not arg.runlocal else 50
+    config.total_steps = 100_000 if not arg.runlocal else 10
+
+  def tokenizer(inkey, outkey='labels'):
+    return (f'render_unifont('
+            f'inkey="{inkey}", '
+            f'outkey="{outkey}", '
+            f'image_size={arg.res}, '
+            f'lower=True, '
+            f'font_size=16, '
+            f'text_brightness=0, '
+            f'background_brightness=127)|'
+            f'value_range(-1, 1, inkey="{outkey}", outkey="{outkey}")')
+
+  pp_image = f'decode|resize({arg.res})|value_range(-1,1)'
+  if arg.test_with_coco:
+    # Train with augmentation when sanity-checking
+    pp_image_aug = (
+        f'decode|resize({arg.res})|flip_lr|randaug(2,10)|value_range(-1,1)')
+    config.input.pp = pp_eval = (
+        f'{pp_image_aug}|flatten|{tokenizer("captions/text")}|'
+        f'keep("image", "labels")')
+  else:
+    config.input.pp = pp_eval = (
+        f'{pp_image}|flatten|{tokenizer("text")}|keep("image", "labels")')
+
+  config.pp_modules = [
+      'ops_general', 'ops_image', 'ops_text', 'proj.clippo.pp_ops']
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = 5000
+
+  config.loss_use_global_batch = True
+
+  # Define the model
+  config.model_name = 'proj.clippo.one_tower'
+
+  config.model = ConfigDict()
+  config.model.image_model = 'vit'
+  config.model.image = ConfigDict({
+      'variant': arg.variant,
+      'pool_type': 'map',
+      'head_zeroinit': False,
+  })
+
+  if arg.test_with_coco:
+    # Initialize with ImageNet21k pretrained checkpoint for sanity-checking
+    assert arg.variant == 'B/16', arg.variant
+    config.model_init = {'image': 'howto-i21k-B/16'}
+    config.model_load = {}
+    config.model_load['img_load_kw'] = {
+        'dont_load': ['^head/.*', '^MAPHead_0/.*', 'cls']}
+
+  config.model.temperature_init = 10.0
+  config.model.out_dim = 768
+
+  # Define the optimizer
+  config.optax_name = 'big_vision.scale_by_adafactor'
+  config.grad_clip_norm = 1.0
+
+  if arg.test_with_coco:
+    # Short schedule for sanity-checking
+    config.lr = 0.0001
+    config.wd = 0.0003
+    config.schedule = dict(decay_type='rsqrt',
+                           timescale=100,
+                           warmup_steps=100 if not arg.runlocal else 5,
+                           cooldown_steps=100 if not arg.runlocal else 5)
+  else:
+    config.lr = 0.001
+    config.wd = 0.0001
+    config.schedule = dict(decay_type='rsqrt',
+                           timescale=10_000,
+                           warmup_steps=10_000 if not arg.runlocal else 5,
+                           cooldown_steps=10_000 if not arg.runlocal else 5)
+
+  # Eval section (Both few-shot and zero-shot)
+  eval_common = dict(
+      type='proj.image_text.contrastive',
+      use_global_batch=config.loss_use_global_batch,
+      log_steps=1000 if not arg.runlocal else 5,
+  )
+  config.evals = {}
+  sub = '[:4]' if arg.runlocal else ''
+  config.evals.val = {
+      **eval_common,
+      'data': val_data,
+      'pp_fn': pp_eval,
+  }
+  config.evals.coco = {
+      **eval_common,
+      'data': dict(name='coco_captions', split=f'val{sub}'),
+      'pp_fn': (
+          f'{pp_image}|flatten|{tokenizer("captions/text")}|'
+          f'keep("image", "labels")'),
+  }
+
+  if arg.i1k_eval:
+    # Requires manual download, see
+    # https://github.com/google-research/big_vision#preparing-tfds-data
+    config.evals.imagenet = {
+        **eval_common,
+        'data': dict(name='imagenet2012', split=f'validation{sub}'),
+        'pp_fn': (
+            f'{pp_image}|clip_i1k_label_names|'
+            f'{tokenizer("labels")}|keep("image", "labels")'),
+    }
+    config.evals.disclf = dict(
+        type='proj.image_text.discriminative_classifier',
+        pp_txt=tokenizer('texts', 'labels'),
+        prefix='z/0shot/',
+        log_steps=5_000 if not arg.runlocal else 5)
+
+  config.evals.retrieval_coco = common.get_coco(
+      pp_img=f'resize({arg.res})|value_range(-1, 1)',
+      pp_txt=tokenizer('texts'),
+      log_steps=5_000 if not arg.runlocal else 5,
+  )
+
+  # Few-shot  metrics
+  config.evals.fewshot = get_fewshot_lsr()
+  config.evals.fewshot.log_steps = 5_000 if not arg.runlocal else 5
+  config.evals.fewshot.representation_layer = 'img/pre_logits'
+
+  config.seed = 0
+
+  return config

big_vision_repo/big_vision/configs/proj/distill/README.md

@@ -0,0 +1,43 @@
+# Knowledge distillation: A good teacher is patient and consistent
+*by Lucas Beyer, Xiaohua Zhai, Amélie Royer, Larisa Markeeva, Rohan Anil, Alexander Kolesnikov*
+
+## Introduction
+We publish all teacher models, and configurations for the main experiments of
+the paper, as well as training logs and student models.
+
+Please read the main [big_vision README](/README.md) to learn how to run
+configs, and remember that each config file contains an example invocation in
+the top-level comment.
+
+## Results
+
+We provide the following [colab to read and plot the logfiles](https://colab.research.google.com/drive/1nMykzUzsfQ_uAxfj3k35DYsATnG_knPl?usp=sharing)
+of a few runs that we reproduced on Cloud.
+
+### ImageNet-1k
+
+The file [bit_i1k.py](bit_i1k.py) is the configuration which reproduces our
+distillation runs on ImageNet-1k reported in Figures 1 and 5(left) and the first
+row of Table1.
+
+We release both student and teacher models:
+
+| Model      | Download link | Resolution  | ImageNet top-1 acc. (paper) | 
+| :---       | :---:         | :---:       |  :---:                      |
+| BiT-R50x1  | [link](https://storage.googleapis.com/bit_models/distill/R50x1_160.npz)      | 160 |  80.5 |
+| BiT-R50x1  | [link](https://storage.googleapis.com/bit_models/distill/R50x1_224.npz)      | 224 |  82.8 |
+| BiT-R152x2 | [link](https://storage.googleapis.com/bit_models/distill/R152x2_T_224.npz)   | 224 |  83.0 |
+| BiT-R152x2 | [link](https://storage.googleapis.com/bit_models/distill/R152x2_T_384.npz)   | 384 |  84.3 |
+
+### Flowers/Pet/Food/Sun
+
+The files [bigsweep_flowers_pet.py](bigsweep_flowers_pet.py) and
+[bigsweep_food_sun.py](bigsweep_food_sun.py) can be used to reproduce the
+distillation runs on these datasets and shown in Figures 3,4,9-12, and Table4.
+
+While our open-source release does not currently support doing hyper-parameter
+sweeps, we still provide an example of the sweeps at the end of the configs
+for reference.
+
+### Teacher models
+Links to all teacher models we used can be found in [common.py](common.py).

big_vision_repo/big_vision/configs/proj/distill/bigsweep_flowers_pet.py

@@ -0,0 +1,164 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Distilling BiT-R152x2 into BiT-R50x1 on Flowers/Pet as in https://arxiv.org/abs/2106.05237
+
+While many epochs are required, this is a small dataset, and thus overall it
+is still fast and possible to run on the relatively small v3-8TPUs (or GPUs).
+
+This configuration contains the recommended settings from Fig3/Tab4 of the
+paper, which can be selected via the fast/medium/long config argument.
+(best settings were selected on a 10% minival)
+
+For Flowers:
+- The `fast` variant takes ~1h10m on a v2-8 TPU.
+  Example logs at gs://big_vision/distill/bit_flowers_fast_06-18_2008/big_vision_metrics.txt
+- The `long` variant takes ~25h on a v3-32 TPU.
+  Example logs at gs://big_vision/distill/bit_flowers_long_06-19_0524/big_vision_metrics.txt
+For Pet:
+- The `fast` variant takes ~28min on a v2-8 TPU.
+  Example logs at gs://big_vision/distill/bit_pet_fast_06-16_2338/big_vision_metrics.txt
+- The `long` variant takes ~11h on a v2-8 and ~8h on a v3-32.
+  Example logs at gs://big_vision/distill/bit_pet_long_06-17_0050/big_vision_metrics.txt
+
+big_vision.trainers.proj.distill.distill \
+    --config big_vision/configs/proj/distill/bigsweep_flowers_pet.py:data=flowers,variant=fast \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+"""
+
+import big_vision.configs.common as bvcc
+import big_vision.configs.proj.distill.common as cd
+import ml_collections as mlc
+
+NCLS = dict(flowers=102, pet=37)
+
+
+def get_config(arg=None):
+  """Config for massive hypothesis-test on pet."""
+  arg = bvcc.parse_arg(arg, runlocal=False, data='flowers', variant='medium', crop='inception_crop(128)')
+  config = mlc.ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(
+      name=dict(flowers='oxford_flowers102', pet='oxford_iiit_pet')[arg.data],
+      split=dict(flowers='train', pet='train[:90%]')[arg.data],
+  )
+  config.input.batch_size = 512
+  config.input.cache_raw = True
+  config.input.shuffle_buffer_size = 50_000
+  config.prefetch_to_device = 4
+
+  config.num_classes = NCLS[arg.data]
+  config.total_epochs = {
+      'flowers': {'fast': 10_000, 'medium': 100_000, 'long': 1_000_000},
+      'pet': {'fast': 1000, 'medium': 3000, 'long': 30_000},
+  }[arg.data][arg.variant]
+
+  config.log_training_steps = 100
+  config.ckpt_steps = 2500
+
+  # Model section
+  config.student_name = 'bit_paper'
+  config.student = dict(depth=50, width=1)
+
+  config.teachers = ['prof_m']
+  config.prof_m_name = 'bit_paper'
+  config.prof_m_init = cd.inits[f'BiT-M R152x2 {arg.data} rc128']
+  config.prof_m = dict(depth=152, width=2)
+
+  # Preprocessing pipeline for student & tacher.
+  pp_common = (
+      '|value_range(-1, 1)'
+      f'|onehot({config.num_classes}, key="label", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+  config.input.pp = f'decode|{arg.crop}|flip_lr' + pp_common
+  ppv = 'decode|resize_small(160)|central_crop(128)' + pp_common
+
+  config.mixup = dict(p=1.0)
+
+  # Distillation settings
+  config.distance = 'kl'
+  config.distance_kw = dict(t={
+      'flowers': {'fast': 10., 'medium': 1., 'long': 1.},
+      'pet': {'fast': 5., 'medium': 10., 'long': 2.},
+  }[arg.data][arg.variant])
+
+  # Optimizer section
+  config.grad_clip_norm = 1.0
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(mu_dtype='bfloat16')
+
+  config.lr = {
+      'flowers': {'fast': 0.003, 'medium': 0.001, 'long': 0.0003},
+      'pet': {'fast': 0.01, 'medium': 0.003, 'long': 0.003},
+  }[arg.data][arg.variant]
+  config.wd = {
+      'flowers': {'fast': 3e-4, 'medium': 1e-4, 'long': 1e-5},
+      'pet': {'fast': 1e-3, 'medium': 3e-4, 'long': 1e-5},
+  }[arg.data][arg.variant]
+  config.schedule = dict(warmup_steps=1500, decay_type='cosine')
+  config.optim_name = 'adam_hp'
+
+  # Eval section
+  minitrain_split = 'train[:512]' if not arg.runlocal else 'train[:16]'
+  if arg.data == 'flowers':
+    val_split = 'validation' if not arg.runlocal else 'validation[:16]'
+    test_split = 'test' if not arg.runlocal else 'test[:16]'
+  elif arg.data == 'pet':
+    val_split = 'train[90%:]' if not arg.runlocal else 'train[:16]'
+    test_split = 'test' if not arg.runlocal else 'test[:16]'
+
+  def get_eval(split):
+    return dict(
+        type='classification',
+        pred='student_fwd',
+        data=dict(name=config.input.data.name, split=split),
+        pp_fn=ppv,
+        loss_name='softmax_xent',
+        log_steps=500,
+    )
+  config.evals = {}
+  config.evals.student_train = get_eval(minitrain_split)
+  config.evals.student_val = get_eval(val_split)
+  config.evals.student_test = get_eval(test_split)
+
+  # Teacher is fixed, so rare evals.
+  teacher = dict(log_steps=100_000, pred='prof_m_fwd')
+  config.evals.teacher_train = {**config.evals.student_train, **teacher}
+  config.evals.teacher_val = {**config.evals.student_val, **teacher}
+  config.evals.teacher_test = {**config.evals.student_test, **teacher}
+
+  # Could in principle also look at agreement on other datasets!
+  def get_dist(split):
+    return dict(
+        type='proj.distill.distance',
+        pred='student_prof_m_fwd',
+        data=dict(name=config.input.data.name, split=split),
+        pp_fn=ppv + '|keep("image")',
+        log_steps=1000,
+        distances=({'kind': 'kl'}, {'kind': 'euclidean'},
+                   {'kind': 'agree', 'k': 1}, {'kind': 'agree', 'k': 5}),
+    )
+  config.evals.dist_train = get_dist(minitrain_split)
+  config.evals.dist_val = get_dist(val_split)
+  config.evals.dist_test = get_dist(test_split)
+
+  # Make a few things much smaller for quick local debugging testruns.
+  if arg.runlocal:
+    config.input.shuffle_buffer_size = 10
+    config.input.batch_size = 8
+
+  return config

big_vision_repo/big_vision/configs/proj/distill/bigsweep_food_sun.py

@@ -0,0 +1,213 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Distilling BiT-R152x2 into BiT-R50x1 on Food101/Sun397 as in https://arxiv.org/abs/2106.05237
+
+While many epochs are required, this is a small dataset, and thus overall it
+is still fast and possible to run on the relatively small v3-8TPUs (or GPUs).
+
+This configuration contains the recommended settings from Fig3/Tab4 of the
+paper, which can be selected via the fast/medium/long config argument.
+(best settings were selected on a 10% minival)
+
+For Food101:
+- The `fast` variant takes ~45min on a v2-8 TPU.
+  Example logs at gs://big_vision/distill/bit_food_fast_06-19_0547/big_vision_metrics.txt
+  Example logs at gs://big_vision/distill/bit_sun_fast_06-20_1839/big_vision_metrics.txt
+- The `long` variant takes ~14h on a v3-8 TPU.
+  Example logs at gs://big_vision/distill/bit_food_long_06-19_0614/big_vision_metrics.txt
+  Example logs at gs://big_vision/distill/bit_sun_long_06-20_1912/big_vision_metrics.txt
+
+big_vision.trainers.proj.distill.distill \
+    --config big_vision/configs/proj/distill/bigsweep_food_sun.py:data=food,variant=fast \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+"""
+
+import big_vision.configs.common as bvcc
+import big_vision.configs.proj.distill.common as cd
+import ml_collections as mlc
+
+H, L = 160, 128
+NCLS = dict(food=101, sun=397)
+
+
+def get_config(arg=None):
+  """Config for massive hypothesis-test on pet."""
+  arg = bvcc.parse_arg(arg, runlocal=False, data='food', variant='medium', crop='inception_crop(128)')
+  config = mlc.ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(
+      name=dict(food='food101', sun='sun397')[arg.data],
+      split=dict(food='train[:90%]', sun='train')[arg.data],
+  )
+  config.input.batch_size = 512
+  config.input.cache_raw = True
+  config.input.shuffle_buffer_size = 50_000
+  config.prefetch_to_device = 4
+
+  config.num_classes = NCLS[arg.data]
+  config.total_epochs = {'fast': 100, 'medium': 1000, 'long': 3000}[arg.variant]
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 2500
+
+  # Model section
+  config.student_name = 'bit_paper'
+  config.student = dict(depth=50, width=1)
+
+  config.teachers = ['prof_m']
+  config.prof_m_name = 'bit_paper'
+  config.prof_m_init = cd.inits[f'BiT-M R152x2 {arg.data} rc128']
+  config.prof_m = dict(depth=152, width=2)
+
+  # Preprocessing pipeline for student & tacher.
+  pp_common = (
+      '|value_range(-1, 1)'
+      f'|onehot({config.num_classes}, key="label", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+  config.input.pp = f'decode|{arg.crop}|flip_lr' + pp_common
+  ppv = 'decode|resize_small(160)|central_crop(128)' + pp_common
+
+  config.mixup = dict(p=1.0)
+
+  # Distillation settings
+  config.distance = 'kl'
+  config.distance_kw = dict(t={
+      'food': {'fast': 10., 'medium': 10., 'long': 5.},
+      'sun': {'fast': 10., 'medium': 10., 'long': 10.},
+  }[arg.data][arg.variant])
+
+  # Optimizer section
+  config.grad_clip_norm = 1.0
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(mu_dtype='bfloat16')
+
+  config.lr = {
+      'food': {'fast': 0.01, 'medium': 0.001, 'long': 0.01},
+      'sun': {'fast': 0.01, 'medium': 0.001, 'long': 0.01},
+  }[arg.data][arg.variant]
+  config.wd = {
+      'food': {'fast': 1e-3, 'medium': 3e-4, 'long': 1e-4},
+      'sun': {'fast': 1e-3, 'medium': 1e-4, 'long': 3e-5},
+  }[arg.data][arg.variant]
+  config.schedule = dict(warmup_steps=1500, decay_type='cosine')
+  config.optim_name = 'adam_hp'
+
+  # Eval section
+  minitrain_split = 'train[:1024]' if not arg.runlocal else 'train[:16]'
+  if arg.data == 'food':
+    val_split = 'train[90%:]' if not arg.runlocal else 'train[:16]'
+    test_split = 'validation' if not arg.runlocal else 'test[:16]'
+  elif arg.data == 'sun':
+    val_split = 'validation' if not arg.runlocal else 'validation[:16]'
+    test_split = 'test' if not arg.runlocal else 'test[:16]'
+
+  def get_eval(split):
+    return dict(
+        type='classification',
+        pred='student_fwd',
+        data=dict(name=config.input.data.name, split=split),
+        pp_fn=ppv,
+        loss_name='softmax_xent',
+        log_steps=500,
+    )
+  config.evals = {}
+  config.evals.student_train = get_eval(minitrain_split)
+  config.evals.student_val = get_eval(val_split)
+  config.evals.student_test = get_eval(test_split)
+
+  # Teacher is fixed, so rare evals.
+  teacher = dict(log_steps=100_000, pred='prof_m_fwd')
+  config.evals.teacher_train = {**config.evals.student_train, **teacher}
+  config.evals.teacher_val = {**config.evals.student_val, **teacher}
+  config.evals.teacher_test = {**config.evals.student_test, **teacher}
+
+  # Could in principle also look at agreement on other datasets!
+  def get_dist(split):
+    return dict(
+        type='proj.distill.distance',
+        pred='student_prof_m_fwd',
+        data=dict(name=config.input.data.name, split=split),
+        pp_fn=ppv + '|keep("image")',
+        log_steps=1000,
+        distances=({'kind': 'kl'}, {'kind': 'euclidean'},
+                   {'kind': 'agree', 'k': 1}, {'kind': 'agree', 'k': 5}),
+    )
+  config.evals.dist_train = get_dist(minitrain_split)
+  config.evals.dist_val = get_dist(val_split)
+  config.evals.dist_test = get_dist(test_split)
+
+  # Make a few things much smaller for quick local debugging testruns.
+  if arg.runlocal:
+    config.input.shuffle_buffer_size = 10
+    config.input.batch_size = 8
+
+  return config
+
+
+def get_hyper(hyper):
+  """Hyper sweep."""
+  # TODO: update, similar to flowers_pet sweep.
+  # By default, not running the MASSIVE sweep, just the recommended setting
+  # across durations. However, code for sweep is left for reference/convenience.
+  return hyper.zipit([
+      hyper.sweep('config.total_epochs', [100, 1_000]),
+      hyper.sweep('config.mixup.p', [0.0, 1.0]),
+      hyper.sweep('config.weight_decay', [1e-3, 1e-5]),
+  ])
+
+  # pylint: disable=unreachable
+
+  def fix(**kw):
+    return hyper.product([hyper.fixed(f'config.{k}', v, length=1)
+                          for k, v in kw.items()])
+
+  def setting(p, l, m, crop, pp_end=None, **extra):
+    pp_end = pp_end or (
+        f'|value_range(-1, 1, key="image")'
+        f'|onehot({NCLS}, key="label", key_result="labels")'
+        f'|keep("image", "labels")'
+    )
+    return hyper.product([
+        fix(**{'mixup.p': p}),
+        fix(l=l, m=m, crop=crop),
+        fix(pp_train=f'decode|{crop}|flip_lr|randaug({l},{m})' + pp_end),
+        fix(**extra)
+    ])
+
+  # Mixup, Layers and Mag in randaug.
+  plm = [(0.0, 0, 0), (0.1, 0, 0), (0.5, 0, 0), (1.0, 0, 0)]
+  return hyper.product([
+      hyper.sweep('config.total_epochs', [100, 1000, 3000]),
+      hyper.sweep('config.lr.base', [0.001, 0.003, 0.01]),
+      hyper.sweep('config.distance_kw.t', [1.0, 2.0, 5.0, 10.0]),
+      hyper.sweep('config.weight_decay', [1e-5, 3e-5, 1e-4, 3e-4, 1e-3]),
+      hyper.chainit(
+          [setting(p=p, l=l, m=m,
+                   crop=(f'resize({H})'
+                         f'|inception_crop({L}, outkey="student")'
+                         f'|central_crop({L}, outkey="teacher")'),
+                   pp_end=(
+                       f'|value_range(-1, 1, key="student")'
+                       f'|value_range(-1, 1, key="teacher")'
+                       f'|onehot({NCLS}, key="label", key_result="labels")'
+                       f'|keep("student", "teacher", "labels")'))
+           for p, l, m in plm] +
+          [setting(p=p, l=l, m=m, crop=f'inception_crop({L})') for
+           p, l, m in plm],
+      )
+  ])

big_vision_repo/big_vision/configs/proj/distill/bit_i1k.py

@@ -0,0 +1,167 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Distilling BiT-R152x2 into BiT-R50x1 on ILSVRC-2012 as in https://arxiv.org/abs/2106.05237
+
+Note that as per paper title, good results require many epochs and thus
+a lot of _patience_. For experimentation/exploration, consider
+using the smaller datasets.
+
+300ep take about 15h on a v3-32 TPU, an example log is available at:
+  Example logs at gs://big_vision/distill/bit_i1k_300ep_06-16/big_vision_metrics.txt
+
+big_vision.trainers.proj.distill.distill \
+    --config big_vision/configs/proj/distill/bit_i1k.py \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+    --config.total_epochs 1200
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+import big_vision.configs.proj.distill.common as cd
+import ml_collections as mlc
+
+
+def get_config(arg=None):
+  """Config for distilling on ImageNet."""
+  arg = bvcc.parse_arg(arg, runlocal=False)
+  config = mlc.ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(name='imagenet2012', split='train[:98%]')
+  config.input.batch_size = 4096
+  config.input.shuffle_buffer_size = 250_000
+
+  config.num_classes = 1000
+  config.total_epochs = 1200  # A good middle-ground
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = 20000
+
+  # Model section
+  config.student_name = 'bit_paper'
+  config.student = dict(depth=50, width=1)
+
+  config.teachers = ['prof_m']  # You could even add multiple.
+
+  # TODO: use public checkpoint name.
+  config.prof_m_name = 'bit_paper'
+  config.prof_m_init = cd.inits['BiT-M R152x2 imagenet2012 ic224']
+  config.prof_m = dict(depth=152, width=2)
+
+  pp_common = (
+      '|value_range(-1, 1)'
+      '|onehot(1000, key="{lbl}", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+  config.input.pp = (
+      'decode_jpeg_and_inception_crop(224)|flip_lr' +
+      pp_common.format(lbl='label')
+  )
+  ppv = 'decode|resize_small(256)|central_crop(224)' + pp_common
+
+  config.mixup = dict(p=1.0)
+
+  # Distillation settings
+  config.distance = 'kl'
+  config.distance_kw = dict(t=1.0)
+
+  # Optimizer section
+  config.grad_clip_norm = 1.0
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(mu_dtype='bfloat16')
+
+  config.lr = 0.03
+  config.wd = 0.0003
+  config.schedule = dict(warmup_steps=5000, decay_type='cosine')
+
+  # Eval section
+  minitrain_split = 'train[:2%]' if not arg.runlocal else 'train[:16]'
+  minival_split = 'train[99%:]' if not arg.runlocal else 'train[:16]'
+  val_split = 'validation' if not arg.runlocal else 'validation[:16]'
+  real_split = 'validation' if not arg.runlocal else 'validation[:16]'
+  v2_split = 'test' if not arg.runlocal else 'test[:16]'
+
+  def get_eval(split, dataset='imagenet2012'):
+    return dict(
+        type='classification',
+        pred='student_fwd',
+        data=dict(name=dataset, split=split),
+        pp_fn=ppv.format(lbl='label'),
+        loss_name='softmax_xent',
+        log_steps=1000,
+    )
+
+  config.evals = {}
+  config.evals.student_train = get_eval(minitrain_split)
+  config.evals.student_minival = get_eval(minival_split)
+  config.evals.student_val = get_eval(val_split)
+  config.evals.student_v2 = get_eval(v2_split, dataset='imagenet_v2')
+  config.evals.student_real = get_eval(real_split, dataset='imagenet2012_real')
+  config.evals.student_real.pp_fn = ppv.format(lbl='real_label')
+
+  config.evals.student_fewshot = get_fewshot_lsr(runlocal=arg.runlocal)
+  config.evals.student_fewshot.pred = 'student_fwd'
+  config.evals.student_fewshot.log_steps = 10_000
+
+  teacher_eval = dict(
+      log_steps=100_000,  # Teacher is fixed, so rare evals.
+      pred='prof_m_fwd',
+  )
+  config.evals.teacher_train = {**config.evals.student_train, **teacher_eval}
+  config.evals.teacher_minival = {**config.evals.student_minival, **teacher_eval}
+  config.evals.teacher_val = {**config.evals.student_val, **teacher_eval}
+  config.evals.teacher_v2 = {**config.evals.student_v2, **teacher_eval}
+  config.evals.teacher_real = {**config.evals.student_real, **teacher_eval}
+  config.evals.teacher_fewshot = {**config.evals.student_fewshot, **teacher_eval}
+  config.evals.teacher_fewshot.prefix = 'z_teacher/'
+
+  # Could in principle also look at agreement on other datasets!
+  def get_dist(split, dataset='imagenet2012'):
+    return dict(
+        type='proj.distill.distance',
+        pred='student_prof_m_fwd',
+        data=dict(name=dataset, split=split),
+        pp_fn=ppv.format(lbl='label') + '|keep("image")',
+        log_steps=1000,
+        distances=({'kind': 'kl'}, {'kind': 'euclidean'},
+                   {'kind': 'agree', 'k': 1}, {'kind': 'agree', 'k': 5}),
+    )
+  config.evals.dist_train = get_dist(minitrain_split)
+  config.evals.dist_minival = get_dist(minival_split)
+  config.evals.dist_val = get_dist(val_split)
+  config.evals.dist_v2 = get_dist(v2_split, dataset='imagenet_v2')
+
+  # NOTE: CKA evaluator does not work with batch padding, so the size of the
+  # split must be a multiple of the batch size.
+  def get_cka(split):
+    return dict(
+        type='proj.distill.cka',
+        pred='student_prof_m_fwd',
+        data=dict(name='imagenet2012', split=split),
+        pp_fn=ppv.format(lbl='label') + '|keep("image")',
+        log_steps=1000,
+    )
+  config.evals.cka_train = get_cka('train[:24576]' if not arg.runlocal else 'train[:16]')
+  config.evals.cka_minival = get_cka('train[-24576:]' if not arg.runlocal else 'train[:16]')
+  config.evals.cka_val = get_cka('validation[:49152]' if not arg.runlocal else 'validation[:16]')
+
+  # Make a few things much smaller for quick local debugging testruns.
+  if arg.runlocal:
+    config.input.shuffle_buffer_size = 10
+    config.input.batch_size = 8
+
+  return config

big_vision_repo/big_vision/configs/proj/distill/common.py

@@ -0,0 +1,27 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Most common teachers for distillation."""
+
+# pylint: disable=line-too-long
+inits = {  # pylint: disable=duplicate-key Internally, we override some paths for convenience.
+    'BiT-M R152x2 imagenet2012 ic224': 'gs://bit_models/distill/R152x2_T_224.npz',
+    'BiT-M R152x2 imagenet2012 rc384': 'gs://bit_models/distill/R152x2_T_384.npz',
+    'BiT-M R152x2 flowers rc128': 'gs://bit_models/distill/R152x2_T_flowers128.npz',
+    'BiT-M R152x2 pet rc128': 'gs://bit_models/distill/R152x2_T_pet128.npz',
+    'BiT-M R152x2 food rc128': 'gs://bit_models/distill/R152x2_T_food128.npz',
+    'BiT-M R152x2 sun rc128': 'gs://bit_models/distill/R152x2_T_sun128.npz',
+
+}
+# pylint: enable=line-too-long

big_vision_repo/big_vision/configs/proj/flexivit/README.md

@@ -0,0 +1,64 @@
+# FlexiViT: One Model for All Patch Sizes
+*by Lucas Beyer, Pavel Izmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, Filip Pavetic*
+
+## Introduction
+We publish all pre-trained FlexiViT models, and configurations for training
+those, as well as training logs for one run.
+
+Please read the main [big_vision README](/README.md) to learn how to run
+configs, and remember that each config file contains an example invocation in
+the top-level comment.
+
+## Pre-trained paper models
+
+Here are the models that we used as backbones in the paper. See Tables in the
+appendix of the paper for expected scores at various patch-sizes and on various
+datasets.
+
+First, the recommended models we used for all experiments.
+Remember that the input is 240px, not 224px:
+
+| Dataset       | Model      | Download link | Notes |
+| :---          | :---:      | :---:         | :---: |
+| ImageNet-1k   | FlexiViT-L | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k.npz) | 1200ep version |
+| ImageNet-1k   | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k.npz) | 1200ep version |
+| ImageNet-1k   | FlexiViT-S | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k.npz) | 1200ep version |
+| ImageNet-21k  | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i21k_300ep.npz) | 300ep version. 1000ep version below is better but was not used in the paper for fair comparison to baselines. |
+| ImageNet-21k  | ViT-B/16   | [link](https://storage.googleapis.com/big_vision/flexivit/vit_b16_i21k_300ep.npz) | Apples-to-apples non-flexi baseline used throughout the paper. |
+| ImageNet-21k  | ViT-B/30   | [link](https://storage.googleapis.com/big_vision/flexivit/vit_b30_i21k_300ep.npz) | Apples-to-apples non-flexi baseline used throughout the paper. |
+
+These models can be used directly in our codebase by specifying
+`model_name = "proj.flexi.vit"` and `model_init = "FlexiViT-L i1k"` for example.
+See the file `models/proj/flexi/vit.py` for more names.
+
+*Important detail:* When further re-using these models with a flexible patch
+size, it is recommended to keep the patch-embedding parameter buffer at its
+original size, and change patch-size on the fly using pi-resize, as opposed to
+changing the parameter buffer's size at load-time.
+For re-using the models with a fixed patch size, either way is fine.
+(The reason is that it is impossible to chain multiple resizes without loss,
+eg doing 32->8->32 does not result in the original weights.)
+
+Second, the list of all released models for completeness:
+
+| Dataset       | Model      | Download link | Notes |
+| :---          | :---:      | :---:         | :---: |
+| ImageNet-21k  | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i21k_1000ep.npz) | 1000ep version. Should be the best available -B model. |
+| ImageNet-21k  | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i21k_90ep.npz) | 90ep version |
+| ImageNet-1k   | FlexiViT-L | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k_600ep.npz) | 600ep version |
+| ImageNet-1k   | FlexiViT-L | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k_300ep.npz) | 300ep version |
+| ImageNet-1k   | FlexiViT-L | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_l_i1k_90ep.npz) | 90ep version |
+| ImageNet-1k   | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k_600ep.npz) | 600ep version |
+| ImageNet-1k   | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k_300ep.npz) | 300ep version |
+| ImageNet-1k   | FlexiViT-B | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_b_i1k_90ep.npz) | 90ep version |
+| ImageNet-1k   | FlexiViT-S | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k_600ep.npz) | 600ep version |
+| ImageNet-1k   | FlexiViT-S | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k_300ep.npz) | 300ep version |
+| ImageNet-1k   | FlexiViT-S | [link](https://storage.googleapis.com/big_vision/flexivit/flexivit_s_i1k_90ep.npz) | 90ep version |
+
+## Results
+
+We provide full training logs for a run with this public code on Cloud that
+reproduces the FlexiViT-S 90ep on i1k results:
+    - [metrics](https://storage.googleapis.com/big_vision/flexivit/deit3_i1k_s_90ep_12-15_2254/big_vision_metrics.txt)
+    - [config](https://storage.googleapis.com/big_vision/flexivit/deit3_i1k_s_90ep_12-15_2254/config.json)
+    - or `gs://big_vision/flexivit/deit3_i1k_s_90ep_12-15_2254`.

big_vision_repo/big_vision/configs/proj/flexivit/i1k_deit3_distill.py

@@ -0,0 +1,187 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Distillation of ViT models into FlexiViT on ImageNet1k.
+
+Run training of the -S variant for 90ep:
+
+big_vision.trainers.proj.flexi.distill \
+  --config big_vision/configs/proj/flexivit/i1k_deit3_distill.py \
+  --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+  --config.total_epochs 90 --config.variant S
+
+Logdir for one reproduction run:
+  - gs://big_vision/flexivit/deit3_i1k_s_90ep_12-15_2254
+
+Timing on Cloud:
+  - S on v3-32: Walltime:10h16m (4h39m eval)
+
+Note that we did not optimize the input for Cloud,
+with tuned caching and prefetching, we should be able to get:
+  - S on v3-32: Walltime: ~6h30m (~1h30m eval)
+  - B on v3-32: Walltime: ~16h00m (~2h30m eval)
+"""
+
+import big_vision.configs.common as bvcc
+
+
+def get_config(arg=None):
+  """Config for distilling ViT on ImageNet1k."""
+  c = bvcc.parse_arg(arg, runlocal=False, res=240)
+
+  c.seed = 0
+  c.total_epochs = 90
+  c.num_classes = 1000
+  c.loss = 'softmax_xent'
+
+  c.input = {}
+  c.input.data = dict(
+      name='imagenet2012',
+      split='train[:99%]',
+  )
+  c.input.batch_size = 1024 if not c.runlocal else 8
+  c.input.cache_raw = False  # Needs up to 120GB of RAM!
+  c.input.shuffle_buffer_size = 250_000 if not c.runlocal else 10
+
+  c.log_training_steps = 50
+  c.ckpt_steps = 1000
+
+  # Model section
+  c.variant = 'B'
+  init = bvcc.format_str('deit3_{variant}_384_1k', c)
+  c.student_name = 'proj.flexi.vit'
+  c.student_init = init
+  c.student = dict(variant=c.get_ref('variant'), pool_type='tok', patch_size=(16, 16))
+
+  c.teachers = ['prof']  # You could even add multiple.
+  c.prof_name = 'vit'
+  c.prof_init = init
+  c.prof = dict(variant=c.get_ref('variant'), pool_type='tok', patch_size=(16, 16))
+
+  pp_label = '|onehot(1000, key="{lbl}", key_result="labels")|keep("image", "prof", "labels")'
+  c.input.pp = (
+      f'decode|inception_crop|flip_lr'
+      '|copy("image", "prof")'
+      f'|resize({c.res})|value_range'
+      '|resize(384, key="prof")|value_range(key="prof")'
+      + pp_label.format(lbl='label')
+  )
+  pp_eval_both = (
+      'decode|copy("image", "prof")|'
+      f'|resize({c.res//7*8})|central_crop({c.res})|value_range'
+      f'|resize({384//7*8}, key="prof")|central_crop(384, key="prof")|value_range(key="prof")|'
+  )
+  pp_eval_student = (
+      f'decode|resize({c.res//7*8})|central_crop({c.res})|value_range(-1, 1)'
+  )
+  pp_eval_prof = (
+      f'decode|resize({384//7*8})|central_crop(384)|value_range(outkey="prof")'
+  )
+
+  c.mixup = dict(p=1.0, n=2)
+
+  # Distillation settings
+  c.distance = 'kl'
+  c.distance_kw = dict(t=1.0)
+
+  # Optimizer section
+  c.grad_clip_norm = 1.0
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(mu_dtype='bfloat16')
+
+  c.lr = 1e-4
+  c.wd = 1e-5
+  c.schedule = dict(warmup_steps=5000, decay_type='cosine')
+
+  # Define the model parameters which are flexible:
+  c.flexi = dict()
+  c.flexi.seqhw = dict(
+      # The settings to sample from. Corresponding patch-sizes at 240px:
+      # 48, 40, 30, 24, 20, 16, 15, 12, 10, 8
+      v=(5, 6, 8, 10, 12, 15, 16, 20, 24, 30),
+      # The probabilities/weights of them. Default uniform.
+      p=(1, 1, 1, 1, 1, 1, 1, 1, 1, 1),
+  )
+
+  # Eval section
+  def mksplit(split):
+    if c.runlocal:
+      return split.split('[')[0] + '[:16]'
+    return split
+
+  minitrain_split = mksplit('train[:2%]')
+  minival_split = mksplit('train[99%:]')
+  val_split = mksplit('validation')
+  test_split = mksplit('test')
+  c.aggressive_cache = False
+
+  def get_eval(s, split, dataset='imagenet2012'):
+    return dict(
+        type='classification',
+        pred=f'student_seqhw={s}',
+        data=dict(name=dataset, split=split),
+        pp_fn=pp_eval_student + pp_label.format(lbl='label'),
+        loss_name='sigmoid_xent',
+        log_percent=0.05,
+        cache_final=False,
+    )
+
+  c.evals = {}
+  for s in c.flexi.seqhw.v:
+    c.evals[f'student_minitrain_{s:02d}'] = get_eval(s, minitrain_split)
+    c.evals[f'student_minival_{s:02d}'] = get_eval(s, minival_split)
+    c.evals[f'student_val_{s:02d}'] = get_eval(s, val_split)
+    c.evals[f'student_v2_{s:02d}'] = get_eval(s, test_split, 'imagenet_v2')
+    c.evals[f'student_a_{s:02d}'] = get_eval(s, test_split, 'imagenet_a')
+    c.evals[f'student_r_{s:02d}'] = get_eval(s, test_split, 'imagenet_r')
+    c.evals[f'student_real_{s:02d}'] = get_eval(s, val_split, 'imagenet2012_real')
+    c.evals[f'student_real_{s:02d}'].pp_fn = pp_eval_student + pp_label.format(lbl='real_label')
+
+  def get_eval_t(split, dataset='imagenet2012'):
+    return dict(
+        type='classification',
+        pred='prof',
+        data=dict(name=dataset, split=split),
+        pp_fn=pp_eval_prof + pp_label.format(lbl='label'),
+        loss_name='sigmoid_xent',
+        log_percent=0.5,  # Teacher is fixed, so eval just for plots.
+        cache_final=False,
+    )
+  c.evals.teacher_minitrain = get_eval_t(minitrain_split)
+  c.evals.teacher_minival = get_eval_t(minival_split)
+  c.evals.teacher_val = get_eval_t(val_split)
+  c.evals.teacher_v2 = get_eval_t(test_split, 'imagenet_v2')
+  c.evals.teacher_a = get_eval_t(test_split, 'imagenet_a')
+  c.evals.teacher_r = get_eval_t(test_split, 'imagenet_r')
+  c.evals.teacher_real = get_eval_t(val_split, 'imagenet2012_real')
+  c.evals.teacher_real.pp_fn = pp_eval_prof + pp_label.format(lbl='real_label')
+
+  # Distance evaluators
+  def get_dist(split, s):
+    return dict(
+        type='proj.distill.distance',
+        pred=f'student_seqhw={s}_prof',
+        data=dict(name='imagenet2012', split=split),
+        pp_fn=pp_eval_both + '|keep("image", "prof")',
+        log_percent=0.05,
+        distances=({'kind': 'kl'}, {'kind': 'logsoftmax_euclidean'},
+                   {'kind': 'agree', 'k': 1}, {'kind': 'agree', 'k': 5}),
+        cache_final=False,
+    )
+  for s in c.flexi.seqhw.v:
+    c.evals[f'dist_minitrain_{s:02d}'] = get_dist(minitrain_split, s)
+    c.evals[f'dist_val_{s:02d}'] = get_dist(val_split, s)
+
+  return c

big_vision_repo/big_vision/configs/proj/flexivit/i21k_distill.py

@@ -0,0 +1,216 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Distill flexible-seqlen ViT on ImageNet-21k from (internal link) B/8.
+
+This config is for reference, we never ran it on public infrastructure.
+
+big_vision.trainers.proj.flexi.distill \
+  --config big_vision/configs/proj/flexivit/i21k_distill.py \
+  --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+  --config.total_epochs 90
+"""
+
+import big_vision.configs.common as bvcc
+
+
+def get_config(arg=None):
+  """Config for training."""
+  # 240px is nice because it's divisible by
+  # [240, 120, 80, 60, 48, 40, 30, 24, 20, 16, 15, 12, 10, 8, 6, 5, 4, 3, 2, 1]
+  c = bvcc.parse_arg(arg, runlocal=False, res=240)
+
+  c.seed = 0
+  c.total_epochs = 90
+  c.num_classes = 21843
+  c.init_head_bias = -10.0
+  c.loss = 'sigmoid_xent'
+
+  c.input = dict()
+  c.input.data = dict(
+      name='imagenet21k',
+      split='full[51200:]',
+  )
+  c.input.batch_size = 4096 if not c.runlocal else 8
+  c.input.shuffle_buffer_size = 250_000 if not c.runlocal else 25
+
+  pp_label_i21k = f'|onehot({c.num_classes})|keep("image", "prof", "labels")'
+  pp_label_i1k = '|onehot(1000, key="{lbl}", key_result="labels")|keep("image", "prof", "labels")'
+  c.input.pp = (
+      f'decode|inception_crop|flip_lr|copy("image", "prof")'
+      f'|resize({c.res})|value_range(-1, 1)'
+      f'|resize(224, outkey="prof")|value_range(-1, 1, key="prof")'
+      + pp_label_i21k
+  )
+  pp_eval_both = (
+      'decode|copy("image", "prof")|'
+      f'|resize_small({c.res//7*8})|central_crop({c.res})|value_range(-1, 1)'
+      f'|resize_small(256, key="prof")|central_crop(224, key="prof")|value_range(-1, 1, key="prof")|'
+  )
+  pp_eval_student = (
+      f'decode|resize({c.res//7*8})|central_crop({c.res})|value_range(-1, 1)'
+  )
+  pp_eval_prof = (
+      'decode|resize(256)|central_crop(224)|value_range(-1, 1, outkey="prof")'
+  )
+
+  # Aggressive pre-fetching because our models here are small, so we not only
+  # can afford it, but we also need it for the smallest models to not be
+  # bottle-necked by the input pipeline. Play around with it for -L models tho.
+  c.input.prefetch = 8
+  c.prefetch_to_device = 4
+
+  c.log_training_steps = 50
+  c.ckpt_steps = 1000
+
+  # Model section
+  init = 'howto-i21k-B/8'
+  c.student_name = 'proj.flexi.vit'
+  c.student_init = init
+  c.student = dict(variant='B', pool_type='tok', patch_size=(8, 8))
+
+  c.teachers = ['prof']  # You could even add multiple.
+  c.prof_name = 'vit'
+  c.prof_init = init
+  c.prof = dict(variant='B/8', pool_type='tok')
+
+  # Define the model parameters which are flexible:
+  c.flexi = dict()
+  c.flexi.seqhw = dict(
+      # The settings to sample from. Corresponding patch-sizes at 240px:
+      # 48, 40, 30, 24, 20, 16, 15, 12, 10, 8
+      v=(5, 6, 8, 10, 12, 15, 16, 20, 24, 30),
+      # The probabilities/weights of them. Default uniform.
+      p=(1, 1, 1, 1, 1, 1, 1, 1, 1, 1),
+  )
+
+  # Distillation settings
+  c.distance = 'kl'
+  c.distance_kw = dict(t=1.0)
+
+  # Optimizer section
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(mu_dtype='bfloat16')
+  c.grad_clip_norm = 1.0
+
+  c.lr = 1e-4
+  c.wd = 1e-5
+  c.schedule = dict(warmup_steps=5000, decay_type='cosine')
+
+  c.mixup = dict(p=1.0)
+
+  ####
+  # Preparing for evals
+  c.evals = {}
+  def mksplit(split):
+    if c.runlocal:
+      return split.split('[')[0] + '[:16]'
+    return split
+
+  ####
+  # Student evals
+
+  # Evaluations on i21k itself.
+  def eval_i21k(s, split):
+    return dict(
+        type='classification',
+        pred=f'student_seqhw={s}',
+        data={**c.input.data, 'split': mksplit(split)},
+        pp_fn=pp_eval_student + pp_label_i21k,
+        loss_name=c.loss,
+        log_steps=5000,  # Very fast O(seconds) so it's fine to run it often.
+    )
+
+  for s in c.flexi.seqhw.v:
+    c.evals[f'student_test{s:02d}'] = eval_i21k(s, 'full[:25_600]')
+    c.evals[f'student_val{s:02d}'] = eval_i21k(s, 'full[25_600:51_200]')
+    c.evals[f'student_minitrain{s:02d}'] = eval_i21k(s, 'full[51_200:76_800]')
+
+  # Evaluations on ImageNet1k variants by label-mapping.
+  def eval_i1k(s, dataset, split, lblmap):
+    return dict(
+        type='classification_with_labelmap',
+        pred=f'student_seqhw={s}',
+        data=dict(name=dataset, split=mksplit(split)),
+        pp_fn=pp_eval_student + pp_label_i1k.format(lbl='label'),
+        loss_name=c.loss,
+        log_steps=5000,  # Very fast O(seconds) so it's fine to run it often.
+        label_mapping=lblmap,
+    )
+  for s in c.flexi.seqhw.v:
+    c.evals[f'student_i1k_val{s:02d}'] = eval_i1k(s, 'imagenet2012', 'validation', 'i1k_i21k')
+    c.evals[f'student_i1k_v2{s:02d}'] = eval_i1k(s, 'imagenet_v2', 'test', 'i1k_i21k')
+    c.evals[f'student_i1k_a{s:02d}'] = eval_i1k(s, 'imagenet_a', 'test', 'i1ka_i21k')
+    c.evals[f'student_i1k_r{s:02d}'] = eval_i1k(s, 'imagenet_r', 'test', 'i1kr_i21k')
+    c.evals[f'student_i1k_real{s:02d}'] = eval_i1k(s, 'imagenet2012_real', 'validation', 'i1k_i21k')
+    c.evals[f'student_i1k_real{s:02d}'].pp_fn = pp_eval_student + pp_label_i1k.format(lbl='real_label')
+    # TODO: add objectnet.
+
+  ####
+  # Teacher evals
+
+  # Evaluations on i21k itself.
+  def eval_i21k_t(split):
+    return dict(
+        type='classification',
+        pred='prof',
+        data={**c.input.data, 'split': mksplit(split)},
+        pp_fn=pp_eval_prof + pp_label_i21k,
+        loss_name=c.loss,
+        log_steps=5000,  # Very fast O(seconds) so it's fine to run it often.
+    )
+
+  c.evals.teacher_test = eval_i21k_t('full[:25_600]')
+  c.evals.teacher_val = eval_i21k_t('full[25_600:51_200]')
+  c.evals.teacher_minitrain = eval_i21k_t('full[51_200:76_800]')
+
+  # Evaluations on ImageNet1k variants by label-mapping.
+  def eval_i1k_t(dataset, split, lblmap):
+    return dict(
+        type='classification_with_labelmap',
+        pred='prof',
+        data=dict(name=dataset, split=mksplit(split)),
+        pp_fn=pp_eval_prof + pp_label_i1k.format(lbl='label'),
+        loss_name=c.loss,
+        log_percent=0.5,  # Teacher is fixed, so eval just for plots.
+        label_mapping=lblmap,
+    )
+  c.evals.teacher_i1k_val = eval_i1k_t('imagenet2012', 'validation', 'i1k_i21k')
+  c.evals.teacher_i1k_v2 = eval_i1k_t('imagenet_v2', 'test', 'i1k_i21k')
+  c.evals.teacher_i1k_a = eval_i1k_t('imagenet_a', 'test', 'i1ka_i21k')
+  c.evals.teacher_i1k_r = eval_i1k_t('imagenet_r', 'test', 'i1kr_i21k')
+  c.evals.teacher_i1k_real = eval_i1k_t('imagenet2012_real', 'validation', 'i1k_i21k')
+  c.evals.teacher_i1k_real.pp_fn = pp_eval_prof + pp_label_i1k.format(lbl='real_label')
+  # TODO: add objectnet.
+
+  ####
+  # Combined evals
+
+  def get_dist(split, s):
+    return dict(
+        type='proj.distill.distance',
+        pred=f'student_seqhw={s}_prof',
+        data=dict(name='imagenet2012', split=mksplit(split)),
+        pp_fn=pp_eval_both + '|keep("image", "prof")',
+        log_percent=0.05,
+        distances=({'kind': 'kl'}, {'kind': 'logsoftmax_euclidean'},
+                   {'kind': 'agree', 'k': 1}, {'kind': 'agree', 'k': 5}),
+    )
+  for s in c.flexi.seqhw.v:
+    c.evals[f'dist_minitrain_{s:02d}'] = get_dist('full[51_200:76_800]', s)
+    c.evals[f'dist_val_{s:02d}'] = get_dist('full[25_600:51_200]', s)
+
+  # Few-shot evaluators not added for overkill reasons for now.
+  return c

big_vision_repo/big_vision/configs/proj/flexivit/i21k_sup.py

@@ -0,0 +1,144 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Pre-training flexible-seqlen ViT on ImageNet-21k following (internal link).
+
+This config is for reference, we never ran it on public infrastructure.
+
+big_vision.trainers.proj.flexi.train \
+  --config big_vision/configs/proj/flexivit/i21k_sup.py \
+  --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+  --config.total_epochs 90
+"""
+
+import big_vision.configs.common as bvcc
+
+
+def get_config(arg=None):
+  """Config for training."""
+  # 240px is nice because it's divisible by
+  # [240, 120, 80, 60, 48, 40, 30, 24, 20, 16, 15, 12, 10, 8, 6, 5, 4, 3, 2, 1]
+  c = bvcc.parse_arg(arg, runlocal=False, res=240)
+
+  c.seed = 0
+  c.total_epochs = 90
+  c.num_classes = 21843
+  c.init_head_bias = -10.0
+  c.loss = 'sigmoid_xent'
+
+  c.input = dict()
+  c.input.data = dict(
+      name='imagenet21k',
+      split='full[51200:]',
+  )
+  c.input.batch_size = 4096 if not c.runlocal else 8
+  c.input.shuffle_buffer_size = 250_000 if not c.runlocal else 25
+
+  pp_common = '|value_range(-1, 1)|onehot({onehot_args})|keep("image", "labels")'
+  pp_common_i21k = pp_common.format(onehot_args=f'{c.num_classes}')
+  pp_common_i1k = pp_common.format(onehot_args='1000, key="{lbl}", key_result="labels"')
+  c.input.pp = f'decode_jpeg_and_inception_crop({c.res})|flip_lr|randaug(2,10)' + pp_common_i21k
+  def pp_eval(res=c.res):
+    return f'decode|resize_small({res//7*8})|central_crop({res})'
+
+  # To continue using the near-defunct randaug op.
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'archive.randaug']
+
+  # Aggressive pre-fetching because our models here are small, so we not only
+  # can afford it, but we also need it for the smallest models to not be
+  # bottle-necked by the input pipeline. Play around with it for -L models tho.
+  c.input.prefetch = 8
+  c.prefetch_to_device = 4
+
+  c.log_training_steps = 50
+  c.ckpt_steps = 1000
+
+  # Model section
+  c.model_name = 'proj.flexi.vit'
+  c.model = dict(
+      variant='B',
+      pool_type='tok',
+      posemb='learn',
+      # patch_size=(32, 32),
+      patch_size=(8, 8),
+      posemb_size=(7, 7),
+      seqhw=None,  # Dynamic!
+  )
+
+  # Define the model parameters which are flexible:
+  c.flexi = dict()
+  c.flexi.seqhw = dict(
+      # The settings to sample from. Corresponding patch-sizes at 240px:
+      # 48, 40, 30, 24, 20, 16, 15, 12, 10, 8
+      v=(5, 6, 8, 10, 12, 15, 16, 20, 24, 30),
+      # The probabilities/weights of them. Default uniform.
+      p=(1, 1, 1, 1, 1, 1, 1, 1, 1, 1),
+  )
+
+  # Optimizer section
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(mu_dtype='bfloat16')
+  c.grad_clip_norm = 1.0
+
+  c.lr = 0.001
+  c.wd = 0.0001
+  c.schedule = dict(warmup_steps=10_000, decay_type='cosine')
+
+  c.mixup = dict(p=0.2, fold_in=None)
+
+  def mksplit(split):
+    if c.runlocal:
+      return split.split('[')[0] + '[:16]'
+    return split
+
+  # Evaluations on i21k itself.
+  def eval_i21k(s, split):
+    return dict(
+        type='classification',
+        pred=f'predict_seqhw={s}',
+        data={**c.input.data, 'split': mksplit(split)},
+        pp_fn=pp_eval() + pp_common_i21k,
+        loss_name=c.loss,
+        log_steps=5000,  # Very fast O(seconds) so it's fine to run it often.
+    )
+
+  c.evals = {}
+  for s in c.flexi.seqhw.v:
+    c.evals[f'test{s:02d}'] = eval_i21k(s, 'full[:25_600]')
+    c.evals[f'val{s:02d}'] = eval_i21k(s, 'full[25_600:51_200]')
+    c.evals[f'train{s:02d}'] = eval_i21k(s, 'full[51_200:76_800]')
+
+  # Evaluations on ImageNet1k variants by label-mapping.
+  def eval_i1k(s, dataset, split, lblmap):
+    return dict(
+        type='classification_with_labelmap',
+        pred=f'predict_seqhw={s}',
+        data=dict(name=dataset, split=mksplit(split)),
+        pp_fn=pp_eval() + pp_common_i1k.format(lbl='label'),
+        loss_name=c.loss,
+        log_steps=5000,  # Very fast O(seconds) so it's fine to run it often.
+        label_mapping=lblmap,
+    )
+  for s in c.flexi.seqhw.v:
+    c.evals[f'i1k_val{s:02d}'] = eval_i1k(s, 'imagenet2012', 'validation', 'i1k_i21k')
+    c.evals[f'i1k_v2{s:02d}'] = eval_i1k(s, 'imagenet_v2', 'test', 'i1k_i21k')
+    c.evals[f'i1k_a{s:02d}'] = eval_i1k(s, 'imagenet_a', 'test', 'i1ka_i21k')
+    c.evals[f'i1k_r{s:02d}'] = eval_i1k(s, 'imagenet_r', 'test', 'i1kr_i21k')
+    c.evals[f'i1k_real{s:02d}'] = eval_i1k(s, 'imagenet2012_real', 'validation', 'i1k_i21k')
+    c.evals[f'i1k_real{s:02d}'].pp_fn = pp_eval() + pp_common_i1k.format(lbl='real_label')
+    # TODO: add objectnet.
+
+  # Few-shot evaluators not added for overkill reasons for now.
+  return c

big_vision_repo/big_vision/configs/proj/flexivit/timing.py

@@ -0,0 +1,53 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long,missing-function-docstring
+r"""A config to run timing for FlexiViT (only inference, no I/O etc.).
+
+big_vision.tools.eval_only \
+  --config big_vision/configs/proj/flexivit/timing.py \
+  --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'` \
+  --config.total_epochs 90
+"""
+
+from ml_collections import ConfigDict
+
+
+def get_config():
+  c = ConfigDict()
+
+  shape = (240, 240, 3)
+  c.batch_size = 8  # swept
+  c.init_shapes = [(1, *shape)]
+  c.representation_layer = 'pre_logits'
+
+  # Creating complete model using all params, the sweep will go over variants.
+  c.model_name = 'xp.flexivit.vit'
+  c.model = dict(
+      variant='B',
+      pool_type='tok',
+      patch_size=(10, 10),  # Like deit@384
+      seqhw=(24, 24),
+  )
+  c.num_classes = 0
+
+  c.evals = {}
+  c.evals.timing = dict(
+      type='timing',
+      input_shapes=[shape],
+      timing=True,
+      pred_kw=dict(outputs=('pre_logits',)),
+  )
+
+  return c

big_vision_repo/big_vision/configs/proj/givt/README.md

@@ -0,0 +1,111 @@
+# GIVT: Generative Infinite-Vocabulary Transformers
+
+*by Michael Tschannen, Cian Eastwood, Fabian Mentzer* [[arxiv]](https://arxiv.org/abs/2312.02116) [[colab]](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/givt/givt_demo_colab.ipynb)
+
+![GIVT overview](givt_overview.png)
+
+
+### Summary
+
+We introduce generative infinite-vocabulary transformers (GIVT) which generate vector sequences with real-valued entries, instead of discrete tokens from a finite vocabulary.
+To this end, we propose two surprisingly simple modifications to decoder-only transformers: 1) at the input, we replace the finite-vocabulary lookup table with a linear projection of the input vectors; and 2) at the output, we replace the logits prediction (usually mapped to a categorical distribution) with the parameters of a multivariate Gaussian mixture model.
+Inspired by the image-generation paradigm of VQ-GAN and MaskGIT, where transformers are used to model the discrete latent sequences of a VQ-VAE, we use GIVT to model the unquantized real-valued latent sequences of a β-VAE.
+In class-conditional image generation GIVT outperforms VQ-GAN (and improved variants thereof) as well as MaskGIT, and achieves performance competitive with recent latent diffusion models.
+Finally, we obtain strong results outside of image generation when applying GIVT to panoptic segmentation and depth estimation with a VAE variant of the UViM framework.
+
+### Checkpoints
+
+We provide model checkpoints for a subset of the models from the paper.
+These are meant as small-scale baselines for researchers interested in exploring GIVT, and are not optimized to provide the best possible visual quality (e.g. scaling the model size can substantially improve visual quality as shown in the paper).
+See below for instructions to train your own models.
+
+**ImageNet 2012 VAEs**
+
+| β     | 1e-5 | 2.5e-5 | 5e-5 | 1e-4 | 2e-4 |
+|:-----------|:------:|:----:|:----:|:----:|:----:|
+| checkpoint | [link][vae_i1k_0] | [link][vae_i1k_1] | [link][vae_i1k_2] | [link][vae_i1k_3] | [link][vae_i1k_4] |
+
+[vae_i1k_0]: https://storage.googleapis.com/big_vision/givt/vae_imagenet_2012_beta_1e-5_params
+[vae_i1k_1]: https://storage.googleapis.com/big_vision/givt/vae_imagenet_2012_beta_2p5e-5_params
+[vae_i1k_2]: https://storage.googleapis.com/big_vision/givt/vae_imagenet_2012_beta_5e-5_params
+[vae_i1k_3]: https://storage.googleapis.com/big_vision/givt/vae_imagenet_2012_beta_1e-4_params
+[vae_i1k_4]: https://storage.googleapis.com/big_vision/givt/vae_imagenet_2012_beta_2e-4_params
+
+**Class-conditional ImageNet 2012 generative models**
+
+| model | resolution | β | inference | FID | checkpoint |
+|:------|:----------:|:------:|:-------------|:---:|:-----------|
+| GIVT-Causal | 256 x 256 | 5e-5 | t=0.95, DB-CFG=0.4 | 3.35 | [link][givt_i1k_1] |
+| GIVT-MaskGIT | 256 x 256 | 5e-5 | t_C=35, DB-CFG=0.1 | 4.53 |  [link][givt_i1k_2] |
+| GIVT-MaskGIT | 512 x 512 | 5e-5 | t_C=140 | 4.86 |  [link][givt_i1k_3] |
+
+[givt_i1k_1]: https://storage.googleapis.com/big_vision/givt/givt_imagenet_2012_causal_params.npz
+[givt_i1k_2]: https://storage.googleapis.com/big_vision/givt/givt_imagenet_2012_maskgit_params.npz
+[givt_i1k_3]: https://storage.googleapis.com/big_vision/givt/givt_imagenet_2012_maskgit_512_params.npz
+
+
+**UViM**
+
+| task | model | dataset | accuracy | checkpoint |
+|:-----|:------|:--------|---------:|:-----------|
+| Panoptic segmentation | VAE (stage 1) | [COCO (2017)] | 71.0 (PQ) | [link][vae_coco_panoptic] |
+| Panoptic segmentation | GIVT (stage 2) | [COCO (2017)] | 40.2 (PQ) | [link][givt_coco_panoptic] |
+| Depth estimation | VAE (stage 1) | [NYU Depth v2] | 0.195 (RMSE) | [link][vae_nyu_depth] |
+| Depth estimation | GIVT (stage 2) | [NYU Depth v2] | 0.474 (RMSE) | [link][givt_nyu_depth] |
+
+[NYU Depth v2]: https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html
+[COCO (2017)]: https://cocodataset.org/#home
+[vae_coco_panoptic]: https://storage.googleapis.com/big_vision/givt/vae_coco_panoptic_params.npz
+[givt_coco_panoptic]: https://storage.googleapis.com/big_vision/givt/givt_coco_panoptic_params.npz
+[vae_nyu_depth]: https://storage.googleapis.com/big_vision/givt/vae_nyu_depth_params.npz
+[givt_nyu_depth]: https://storage.googleapis.com/big_vision/givt/givt_nyu_depth_params.npz
+
+### Training models
+
+This directory contains configs to train GIVT models as well as VAEs (for the UViM variants).
+For training the ImageNet 2012 VAE models we used a modified version of the [MaskGIT code](https://github.com/google-research/maskgit).
+
+The `big_vision` input pipeline relies on [TensorFlow Datasets (TFDS)](https://www.tensorflow.org/datasets) 
+which supports some data sets out-of-the-box, whereas others require manual download of the data
+(for example ImageNet and COCO (2017), see the `big_vision` [main README](../../../../#cloud-tpu-vm-setup) and the [UViM README](../uvim), respectively, for details).
+
+After setting up `big_vision` as described in the [main README](../../../../#cloud-tpu-vm-setup), training can be launched locally as follows
+
+```
+python -m big_vision.trainers.proj.givt.generative \
+  --config big_vision/configs/proj/givt/givt_imagenet2012.py \
+  --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`
+```
+
+Add the suffix `:key1=value1,key2=value2,...` to the config path in the launch
+command to modify the config with predefined arguments (see config for details). For example:
+`--config big_vision/configs/proj/givt/givt_imagenet_2012.py:model_size=large`.
+Note that `givt_imagenet2012.py` uses [Imagenette](https://github.com/fastai/imagenette) to ensure that the config is runnable without manual ImageNet download.
+This is only meant for testing and will overfit immediately. Please download ImageNet to reproduce the paper results.
+
+VAE trainings for the GIVT variant of UViM can be launched as
+
+```
+python -m big_vision.trainers.proj.givt.vae \
+  --config big_vision/configs/proj/givt/vae_nyu_depth.py \
+  --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`
+```
+
+Please refer to the [main README](../../../../#cloud-tpu-vm-setup)
+for details on how to launch training on a (multi-host) TPU setup.
+
+
+### Disclaimer
+
+This is not an official Google Product.
+
+
+### Citation
+```
+@article{tschannen2023givt,
+  title={GIVT: Generative Infinite-Vocabulary Transformers},
+  author={Tschannen, Michael and Eastwood, Cian and Mentzer, Fabian},
+  journal={arXiv:2312.02116},
+  year={2023}
+}
+```

big_vision_repo/big_vision/configs/proj/givt/givt_coco_panoptic.py

@@ -0,0 +1,186 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Train a GIVT encoder-decoder model on COCO panoptic."""
+
+import itertools
+import ml_collections
+
+ConfigDict = ml_collections.ConfigDict
+
+VTT_MODELS = {
+    'base': dict(num_layers=12, num_decoder_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768),
+    'large': dict(num_layers=24, num_decoder_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024),
+}
+
+RES = 512
+PATCH_SIZE = 16
+LABEL_RES = 512
+LABEL_PATCH_SIZE = 16
+
+
+def get_config(runlocal=False):
+  """Config for training."""
+  config = ConfigDict()
+
+  config.input = {}
+  config.input.pp = (
+      f'decode|coco_panoptic|concat(["semantics","instances"], "labels")|'
+      f'randu("fliplr")|det_fliplr(key="image")|det_fliplr(key="labels")|'
+      f'inception_box|crop_box(key="image")|crop_box(key="labels")|'
+      f'resize({RES})|resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1, 1)|make_canonical|'
+      f'copy("image", "cond_image")|copy("labels", "image")|'
+      f'keep("image", "cond_image")'
+  )
+  pp_eval = (
+      f'decode|coco_panoptic|concat(["semantics","instances"], "labels")|'
+      f'resize({RES})|resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1, 1)|make_canonical|'
+      f'copy("image", "cond_image")|copy("labels", "image")|'
+      f'keep("image", "cond_image")'
+  )
+  pp_predict = (
+      f'decode|resize({RES})|value_range(-1, 1)|copy("image", "cond_image")|'
+      f'keep("cond_image", "image/id")'  # image/id used for rng seeds.
+  )
+
+  config.input.data = dict(name='coco/2017_panoptic', split='train[4096:]')
+  config.input.batch_size = 512
+  config.input.shuffle_buffer_size = 50_000
+
+  config.total_epochs = 200
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None
+  config.prefetch_to_device = 2
+  config.seed = 0
+
+  # Optimizer section
+  config.optax_name = 'big_vision.scale_by_adafactor'
+  config.optax = dict(beta2_cap=0.95)
+
+  config.ar_generation_config = ml_collections.ConfigDict()
+  config.ar_generation_config.temp = 0.85
+  config.ar_generation_config.temp_probs = 1.0
+  config.ar_generation_config.beam_size = 4
+  config.ar_generation_config.fan_size = 8
+  config.ar_generation_config.rand_top_k = False
+  config.ar_generation_config.rand_top_k_temp = 1.0
+
+  config.lr = 0.001
+  config.wd = 0.000001
+  config.lr_mults = [
+      ('pos_embedding_encoder.*', 0.1),
+      ('EmbedPatches.*', 0.1),
+      ('encoder.*', 0.1),
+      ('decoder.*', 1.0)
+  ]
+  config.schedule = dict(decay_type='cosine', warmup_steps=4_000)
+
+  # Oracle section
+  config.vae = ConfigDict()
+  config.vae.model_name = 'proj.givt.vit'
+  config.vae.model = ConfigDict()
+  config.vae.model.input_size = (RES, RES)
+  config.vae.model.patch_size = (PATCH_SIZE, PATCH_SIZE)
+  config.vae.model.code_len = 256
+  config.vae.model.width = 768
+  config.vae.model.enc_depth = 6
+  config.vae.model.dec_depth = 12
+  config.vae.model.mlp_dim = 3072
+  config.vae.model.num_heads = 12
+  config.vae.model.codeword_dim = 16
+  config.vae.model.code_dropout = 'none'
+  config.vae.model.bottleneck_resize = True
+  # values: (channel index in source image, number of classes)
+  config.vae.model.inout_specs = {
+      'semantics': (0, 133 + 1),  # +1 for void label
+      'instances': (1, 100),  # COCO: actually 98 train/78 validation.
+  }
+  config.vae.model_init = 'gs://big_vision/givt/vae_coco_panoptic_params.npz'
+
+  # Model section
+  config.model_name = 'proj.givt.givt'
+  # # Base model (for exploration)
+  # config.model_init = {'encoder': 'howto-i21k-B/16'}
+  # config.model = ConfigDict(VTT_MODELS['base'])
+  # Large model
+  config.model_init = {'encoder': 'howto-i21k-L/16'}
+  config.model_load = dict(dont_load=('cls', 'head/bias', 'head/kernel'))
+  config.model = ConfigDict(VTT_MODELS['large'])
+  config.model.patches = (PATCH_SIZE, PATCH_SIZE)
+  config.model.input_size = (RES, RES)
+  config.model.posemb_type = 'learn'
+  config.model.seq_len = config.vae.model.code_len
+  config.model.num_labels = None
+  config.model.num_mixtures = 1
+  config.model.fix_square_plus = True
+  config.model.out_dim = config.vae.model.codeword_dim
+  config.model.scale_tol = 1e-6
+  config.model.dec_dropout_rate = 0.0
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = ConfigDict()
+  config.evals.val.type = 'mean'
+  config.evals.val.pred = 'validation'
+  config.evals.val.data = dict(name=config.input.data.name, split='train[:4096]')
+  config.evals.val.pp_fn = pp_eval
+  config.evals.val.log_steps = 1000
+
+  config.eval_only = False
+
+  base = {
+      'type': 'proj.givt.coco_panoptic',
+      'data': {**config.input.data},
+      'pp_fn': pp_predict,
+      'log_steps': 10_000,
+      'pred': 'sample_panoptic',
+      # Filters objects that occupy less than 0.03^2 fraction of all pixels.
+      # 'pred_kw': {'min_fraction': 0.03 ** 2},
+  }
+  config.evals.coco_panoptic_train = dict(base)
+  config.evals.coco_panoptic_train.data.split = 'train[4096:8192]'
+  config.evals.coco_panoptic_holdout = dict(base)
+  config.evals.coco_panoptic_holdout.data.split = 'train[:4096]'
+  config.evals.coco_panoptic = dict(base)
+  config.evals.coco_panoptic.data.split = 'validation'
+
+  config.evals.save_pred = dict(type='proj.givt.save_predictions')
+  config.evals.save_pred.pred = 'sample_panoptic'
+  config.evals.save_pred.pp_fn = pp_eval
+  config.evals.save_pred.log_steps = 100_000
+  config.evals.save_pred.data = dict(config.input.data)
+  config.evals.save_pred.data.split = 'validation[:1024]'
+  config.evals.save_pred.outfile = 'inference.npz'
+
+  if runlocal:
+    config.input.batch_size = 4
+    config.input.shuffle_buffer_size = 10
+    config.evals.val.data.split = 'train[:16]'
+    config.evals.val.log_steps = 20
+    config.model.num_layers = 1
+    config.model.num_decoder_layers = 1
+    del config.model_init
+    config.evals.val.data.split = 'validation[:4]'
+    config.evals.coco_panoptic.data.split = 'validation[:4]'
+    config.evals.save_pred.data.split = 'validation[:4]'
+    for k in config.evals.keys():
+      if k not in ['val', 'coco_panoptic', 'save_pred']:
+        del config.evals[k]
+
+  return config

big_vision_repo/big_vision/configs/proj/givt/givt_demo_colab.ipynb

@@ -0,0 +1,309 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# GIVT Demo colab\n",
+        "\n",
+        "[[paper]](https://arxiv.org/abs/2312.02116) [[GitHub]](https://github.com/google-research/big_vision/blob/main/big_vision/configs/proj/givt/README.md)\n",
+        "\n",
+        "This colab implements class-conditional image generation using GIVT-Causal and GIVT-MaskGIT for the 1k ImageNet2012 classes.\n",
+        "\n",
+        "The available model checkpoints are meant as small-scale baselines (~300M parameters) for researchers interested in exploring GIVT, and are not optimized to provide the best possible visual quality (e.g. scaling the model size can substantially improve visual quality as shown in the paper).\n",
+        "\n",
+        "The colab was tested with the CPU and T4 GPU runtimes. We recommend the T4 GPU runtime (the CPU rutime is very slow).\n",
+        "\n",
+        "_Disclaimer: This is not an official Google Product._"
+      ],
+      "metadata": {
+        "id": "botgo-GZiWI_"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### `big_vision` setup"
+      ],
+      "metadata": {
+        "id": "jQxc9UZ-mVrQ"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown Clone and set up repository\n",
+        "!git clone --branch=main --depth=1 https://github.com/google-research/big_vision\n",
+        "!cd big_vision && git pull\n",
+        "\n",
+        "# Install dependencies - pin TensorFlow-related packages to ensure compatibility\n",
+        "# which might not be needed in in the future\n",
+        "!echo -e \"keras==3.0.5\\ntensorflow==2.16.1\\ntensorflow-probability==0.24.0\" > big_vision/big_vision/constraints.txt\n",
+        "!pip install -r big_vision/big_vision/requirements.txt -c big_vision/big_vision/constraints.txt\n",
+        "%cd big_vision"
+      ],
+      "metadata": {
+        "id": "ZAXiVta3n2jL",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qYS7JNups4MU",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@markdown Imports\n",
+        "import jax\n",
+        "from functools import partial\n",
+        "import ml_collections\n",
+        "import matplotlib.pyplot as plt\n",
+        "\n",
+        "from big_vision.configs.proj.givt import givt_imagenet2012\n",
+        "from big_vision.datasets.imagenet import class_names as imagenet_class_names\n",
+        "from big_vision.models.proj.givt import givt, cnn, decode, parallel_decode\n",
+        "\n",
+        "jnp = jax.numpy"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Select and download model\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "MaCM_PIcd2Rb"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model = \"GIVT-Causal 256x256\" #@param [\"GIVT-Causal 256x256\", \"GIVT-MaskGIT 256x256\", \"GIVT-MaskGIT 512x512\"]\n",
+        "\n",
+        "givt_ckpt_path, cfg_w, temp, is_ar, res = {\n",
+        "    \"GIVT-Causal 256x256\": (\n",
+        "        \"gs://big_vision/givt/givt_imagenet_2012_causal_params.npz\", 0.4, 0.95, True, 256),\n",
+        "    \"GIVT-MaskGIT 256x256\": (\n",
+        "        \"gs://big_vision/givt/givt_imagenet_2012_maskgit_params.npz\", 0.0, 35.0, False, 256),\n",
+        "    \"GIVT-MaskGIT 512x512\": (\n",
+        "        \"gs://big_vision/givt/givt_imagenet_2012_maskgit_512_params.npz\", 0.0, 140.0, False, 512),\n",
+        "}[model]\n",
+        "\n",
+        "config = givt_imagenet2012.get_config(arg=f\"res={res},style={'ar' if is_ar else 'masked'}\")\n",
+        "\n",
+        "print(\"Loading VAE model...\")\n",
+        "vae_model = cnn.Model(**config.vae.model)\n",
+        "vae_params = cnn.load(None, config.vae.model_init, **config.vae.model_load)\n",
+        "\n",
+        "print(\"Loading GIVT model...\")\n",
+        "givt_model = givt.Model(**config.model)\n",
+        "givt_params = jax.device_put(\n",
+        "    givt.load(None, givt_ckpt_path), jax.devices()[0])"
+      ],
+      "metadata": {
+        "id": "7l6QIjdyN3dg",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### VAE encode/decode and sampling loop"
+      ],
+      "metadata": {
+        "id": "SUj5k1bxd6wr"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "@jax.jit\n",
+        "def vae_encode(images, rng):\n",
+        "  \"\"\"Encode image with VAE encoder.\"\"\"\n",
+        "  mu, logvar = vae_model.apply(\n",
+        "      {\"params\": vae_params}, images, method=vae_model.encode,\n",
+        "  )\n",
+        "  return vae_model.apply(\n",
+        "      {\"params\": vae_params},\n",
+        "      mu,\n",
+        "      logvar,\n",
+        "      method=vae_model.reparametrize,\n",
+        "      rngs={\"dropout\": rng},\n",
+        "  )\n",
+        "\n",
+        "@jax.jit\n",
+        "def vae_decode(z):\n",
+        "  \"\"\"Reconstruct image with VAE decoder from latent code z.\"\"\"\n",
+        "  return vae_model.apply({\"params\": vae_params}, z, method=vae_model.decode)\n",
+        "\n",
+        "### jit-compilation seems to go OOM (RAM) on the free tier GPU colab, but might\n",
+        "### lead to speedups on machines with more resources\n",
+        "# @partial(jax.jit, static_argnums=(2, 3))\n",
+        "def sample(labels, rng, ar_generation_config=None, masked_generation_config=None):\n",
+        "  \"\"\"Sample from GIVT-Causal or GIVT-MaskGIT.\"\"\"\n",
+        "  print(f\"Sampling, style={givt_model.style}\")\n",
+        "  shared_kwargs = dict(\n",
+        "    labels=labels,\n",
+        "    model=givt_model,\n",
+        "    seq_len=config.model.seq_len,\n",
+        "    feature_dim=config.model.out_dim,\n",
+        "  )\n",
+        "\n",
+        "  match givt_model.style:\n",
+        "    case \"ar\":\n",
+        "      sampled_codes, _ = decode.generate(\n",
+        "          params={\"params\": givt_params},\n",
+        "          seed=rng,\n",
+        "          config=dict(ar_generation_config),\n",
+        "          **shared_kwargs,\n",
+        "      )\n",
+        "      info = sampled_codes\n",
+        "    case \"masked\":\n",
+        "      masked_out = parallel_decode.decode_masked(\n",
+        "          rng=rng,\n",
+        "          variables={\"params\": givt_params},\n",
+        "          config=masked_generation_config,\n",
+        "          **shared_kwargs,\n",
+        "      )\n",
+        "      sampled_codes = masked_out.current_inputs_q\n",
+        "      info = masked_out\n",
+        "    case _:\n",
+        "      raise NotImplementedError\n",
+        "  return sampled_codes, info"
+      ],
+      "metadata": {
+        "id": "vSn7Si2FS1zi"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Generate images for class label"
+      ],
+      "metadata": {
+        "id": "tOnWaJZVeOIX"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "rng = 0 #@param = 'int'\n",
+        "label = 'goldfish' #@param [\"tench\", \"goldfish\", \"great white shark\", \"tiger shark\", \"hammerhead shark\", \"electric ray\", \"stingray\", \"rooster\", \"hen\", \"ostrich\", \"brambling\", \"goldfinch\", \"house finch\", \"junco\", \"indigo bunting\", \"American robin\", \"bulbul\", \"jay\", \"magpie\", \"chickadee\", \"American dipper\", \"kite (bird of prey)\", \"bald eagle\", \"vulture\", \"great grey owl\", \"fire salamander\", \"smooth newt\", \"newt\", \"spotted salamander\", \"axolotl\", \"American bullfrog\", \"tree frog\", \"tailed frog\", \"loggerhead sea turtle\", \"leatherback sea turtle\", \"mud turtle\", \"terrapin\", \"box turtle\", \"banded gecko\", \"green iguana\", \"Carolina anole\", \"desert grassland whiptail lizard\", \"agama\", \"frilled-necked lizard\", \"alligator lizard\", \"Gila monster\", \"European green lizard\", \"chameleon\", \"Komodo dragon\", \"Nile crocodile\", \"American alligator\", \"triceratops\", \"worm snake\", \"ring-necked snake\", \"eastern hog-nosed snake\", \"smooth green snake\", \"kingsnake\", \"garter snake\", \"water snake\", \"vine snake\", \"night snake\", \"boa constrictor\", \"African rock python\", \"Indian cobra\", \"green mamba\", \"sea snake\", \"Saharan horned viper\", \"eastern diamondback rattlesnake\", \"sidewinder rattlesnake\", \"trilobite\", \"harvestman\", \"scorpion\", \"yellow garden spider\", \"barn spider\", \"European garden spider\", \"southern black widow\", \"tarantula\", \"wolf spider\", \"tick\", \"centipede\", \"black grouse\", \"ptarmigan\", \"ruffed grouse\", \"prairie grouse\", \"peafowl\", \"quail\", \"partridge\", \"african grey parrot\", \"macaw\", \"sulphur-crested cockatoo\", \"lorikeet\", \"coucal\", \"bee eater\", \"hornbill\", \"hummingbird\", \"jacamar\", \"toucan\", \"duck\", \"red-breasted merganser\", \"goose\", \"black swan\", \"tusker\", \"echidna\", \"platypus\", \"wallaby\", \"koala\", \"wombat\", \"jellyfish\", \"sea anemone\", \"brain coral\", \"flatworm\", \"nematode\", \"conch\", \"snail\", \"slug\", \"sea slug\", \"chiton\", \"chambered nautilus\", \"Dungeness crab\", \"rock crab\", \"fiddler crab\", \"red king crab\", \"American lobster\", \"spiny lobster\", \"crayfish\", \"hermit crab\", \"isopod\", \"white stork\", \"black stork\", \"spoonbill\", \"flamingo\", \"little blue heron\", \"great egret\", \"bittern bird\", \"crane bird\", \"limpkin\", \"common gallinule\", \"American coot\", \"bustard\", \"ruddy turnstone\", \"dunlin\", \"common redshank\", \"dowitcher\", \"oystercatcher\", \"pelican\", \"king penguin\", \"albatross\", \"grey whale\", \"killer whale\", \"dugong\", \"sea lion\", \"Chihuahua\", \"Japanese Chin\", \"Maltese\", \"Pekingese\", \"Shih Tzu\", \"King Charles Spaniel\", \"Papillon\", \"toy terrier\", \"Rhodesian Ridgeback\", \"Afghan Hound\", \"Basset Hound\", \"Beagle\", \"Bloodhound\", \"Bluetick Coonhound\", \"Black and Tan Coonhound\", \"Treeing Walker Coonhound\", \"English foxhound\", \"Redbone Coonhound\", \"borzoi\", \"Irish Wolfhound\", \"Italian Greyhound\", \"Whippet\", \"Ibizan Hound\", \"Norwegian Elkhound\", \"Otterhound\", \"Saluki\", \"Scottish Deerhound\", \"Weimaraner\", \"Staffordshire Bull Terrier\", \"American Staffordshire Terrier\", \"Bedlington Terrier\", \"Border Terrier\", \"Kerry Blue Terrier\", \"Irish Terrier\", \"Norfolk Terrier\", \"Norwich Terrier\", \"Yorkshire Terrier\", \"Wire Fox Terrier\", \"Lakeland Terrier\", \"Sealyham Terrier\", \"Airedale Terrier\", \"Cairn Terrier\", \"Australian Terrier\", \"Dandie Dinmont Terrier\", \"Boston Terrier\", \"Miniature Schnauzer\", \"Giant Schnauzer\", \"Standard Schnauzer\", \"Scottish Terrier\", \"Tibetan Terrier\", \"Australian Silky Terrier\", \"Soft-coated Wheaten Terrier\", \"West Highland White Terrier\", \"Lhasa Apso\", \"Flat-Coated Retriever\", \"Curly-coated Retriever\", \"Golden Retriever\", \"Labrador Retriever\", \"Chesapeake Bay Retriever\", \"German Shorthaired Pointer\", \"Vizsla\", \"English Setter\", \"Irish Setter\", \"Gordon Setter\", \"Brittany dog\", \"Clumber Spaniel\", \"English Springer Spaniel\", \"Welsh Springer Spaniel\", \"Cocker Spaniel\", \"Sussex Spaniel\", \"Irish Water Spaniel\", \"Kuvasz\", \"Schipperke\", \"Groenendael dog\", \"Malinois\", \"Briard\", \"Australian Kelpie\", \"Komondor\", \"Old English Sheepdog\", \"Shetland Sheepdog\", \"collie\", \"Border Collie\", \"Bouvier des Flandres dog\", \"Rottweiler\", \"German Shepherd Dog\", \"Dobermann\", \"Miniature Pinscher\", \"Greater Swiss Mountain Dog\", \"Bernese Mountain Dog\", \"Appenzeller Sennenhund\", \"Entlebucher Sennenhund\", \"Boxer\", \"Bullmastiff\", \"Tibetan Mastiff\", \"French Bulldog\", \"Great Dane\", \"St. Bernard\", \"husky\", \"Alaskan Malamute\", \"Siberian Husky\", \"Dalmatian\", \"Affenpinscher\", \"Basenji\", \"pug\", \"Leonberger\", \"Newfoundland dog\", \"Great Pyrenees dog\", \"Samoyed\", \"Pomeranian\", \"Chow Chow\", \"Keeshond\", \"brussels griffon\", \"Pembroke Welsh Corgi\", \"Cardigan Welsh Corgi\", \"Toy Poodle\", \"Miniature Poodle\", \"Standard Poodle\", \"Mexican hairless dog (xoloitzcuintli)\", \"grey wolf\", \"Alaskan tundra wolf\", \"red wolf or maned wolf\", \"coyote\", \"dingo\", \"dhole\", \"African wild dog\", \"hyena\", \"red fox\", \"kit fox\", \"Arctic fox\", \"grey fox\", \"tabby cat\", \"tiger cat\", \"Persian cat\", \"Siamese cat\", \"Egyptian Mau\", \"cougar\", \"lynx\", \"leopard\", \"snow leopard\", \"jaguar\", \"lion\", \"tiger\", \"cheetah\", \"brown bear\", \"American black bear\", \"polar bear\", \"sloth bear\", \"mongoose\", \"meerkat\", \"tiger beetle\", \"ladybug\", \"ground beetle\", \"longhorn beetle\", \"leaf beetle\", \"dung beetle\", \"rhinoceros beetle\", \"weevil\", \"fly\", \"bee\", \"ant\", \"grasshopper\", \"cricket insect\", \"stick insect\", \"cockroach\", \"praying mantis\", \"cicada\", \"leafhopper\", \"lacewing\", \"dragonfly\", \"damselfly\", \"red admiral butterfly\", \"ringlet butterfly\", \"monarch butterfly\", \"small white butterfly\", \"sulphur butterfly\", \"gossamer-winged butterfly\", \"starfish\", \"sea urchin\", \"sea cucumber\", \"cottontail rabbit\", \"hare\", \"Angora rabbit\", \"hamster\", \"porcupine\", \"fox squirrel\", \"marmot\", \"beaver\", \"guinea pig\", \"common sorrel horse\", \"zebra\", \"pig\", \"wild boar\", \"warthog\", \"hippopotamus\", \"ox\", \"water buffalo\", \"bison\", \"ram (adult male sheep)\", \"bighorn sheep\", \"Alpine ibex\", \"hartebeest\", \"impala (antelope)\", \"gazelle\", \"arabian camel\", \"llama\", \"weasel\", \"mink\", \"European polecat\", \"black-footed ferret\", \"otter\", \"skunk\", \"badger\", \"armadillo\", \"three-toed sloth\", \"orangutan\", \"gorilla\", \"chimpanzee\", \"gibbon\", \"siamang\", \"guenon\", \"patas monkey\", \"baboon\", \"macaque\", \"langur\", \"black-and-white colobus\", \"proboscis monkey\", \"marmoset\", \"white-headed capuchin\", \"howler monkey\", \"titi monkey\", \"Geoffroy's spider monkey\", \"common squirrel monkey\", \"ring-tailed lemur\", \"indri\", \"Asian elephant\", \"African bush elephant\", \"red panda\", \"giant panda\", \"snoek fish\", \"eel\", \"silver salmon\", \"rock beauty fish\", \"clownfish\", \"sturgeon\", \"gar fish\", \"lionfish\", \"pufferfish\", \"abacus\", \"abaya\", \"academic gown\", \"accordion\", \"acoustic guitar\", \"aircraft carrier\", \"airliner\", \"airship\", \"altar\", \"ambulance\", \"amphibious vehicle\", \"analog clock\", \"apiary\", \"apron\", \"trash can\", \"assault rifle\", \"backpack\", \"bakery\", \"balance beam\", \"balloon\", \"ballpoint pen\", \"Band-Aid\", \"banjo\", \"baluster / handrail\", \"barbell\", \"barber chair\", \"barbershop\", \"barn\", \"barometer\", \"barrel\", \"wheelbarrow\", \"baseball\", \"basketball\", \"bassinet\", \"bassoon\", \"swimming cap\", \"bath towel\", \"bathtub\", \"station wagon\", \"lighthouse\", \"beaker\", \"military hat (bearskin or shako)\", \"beer bottle\", \"beer glass\", \"bell tower\", \"baby bib\", \"tandem bicycle\", \"bikini\", \"ring binder\", \"binoculars\", \"birdhouse\", \"boathouse\", \"bobsleigh\", \"bolo tie\", \"poke bonnet\", \"bookcase\", \"bookstore\", \"bottle cap\", \"hunting bow\", \"bow tie\", \"brass memorial plaque\", \"bra\", \"breakwater\", \"breastplate\", \"broom\", \"bucket\", \"buckle\", \"bulletproof vest\", \"high-speed train\", \"butcher shop\", \"taxicab\", \"cauldron\", \"candle\", \"cannon\", \"canoe\", \"can opener\", \"cardigan\", \"car mirror\", \"carousel\", \"tool kit\", \"cardboard box / carton\", \"car wheel\", \"automated teller machine\", \"cassette\", \"cassette player\", \"castle\", \"catamaran\", \"CD player\", \"cello\", \"mobile phone\", \"chain\", \"chain-link fence\", \"chain mail\", \"chainsaw\", \"storage chest\", \"chiffonier\", \"bell or wind chime\", \"china cabinet\", \"Christmas stocking\", \"church\", \"movie theater\", \"cleaver\", \"cliff dwelling\", \"cloak\", \"clogs\", \"cocktail shaker\", \"coffee mug\", \"coffeemaker\", \"spiral or coil\", \"combination lock\", \"computer keyboard\", \"candy store\", \"container ship\", \"convertible\", \"corkscrew\", \"cornet\", \"cowboy boot\", \"cowboy hat\", \"cradle\", \"construction crane\", \"crash helmet\", \"crate\", \"infant bed\", \"Crock Pot\", \"croquet ball\", \"crutch\", \"cuirass\", \"dam\", \"desk\", \"desktop computer\", \"rotary dial telephone\", \"diaper\", \"digital clock\", \"digital watch\", \"dining table\", \"dishcloth\", \"dishwasher\", \"disc brake\", \"dock\", \"dog sled\", \"dome\", \"doormat\", \"drilling rig\", \"drum\", \"drumstick\", \"dumbbell\", \"Dutch oven\", \"electric fan\", \"electric guitar\", \"electric locomotive\", \"entertainment center\", \"envelope\", \"espresso machine\", \"face powder\", \"feather boa\", \"filing cabinet\", \"fireboat\", \"fire truck\", \"fire screen\", \"flagpole\", \"flute\", \"folding chair\", \"football helmet\", \"forklift\", \"fountain\", \"fountain pen\", \"four-poster bed\", \"freight car\", \"French horn\", \"frying pan\", \"fur coat\", \"garbage truck\", \"gas mask or respirator\", \"gas pump\", \"goblet\", \"go-kart\", \"golf ball\", \"golf cart\", \"gondola\", \"gong\", \"gown\", \"grand piano\", \"greenhouse\", \"radiator grille\", \"grocery store\", \"guillotine\", \"hair clip\", \"hair spray\", \"half-track\", \"hammer\", \"hamper\", \"hair dryer\", \"hand-held computer\", \"handkerchief\", \"hard disk drive\", \"harmonica\", \"harp\", \"combine harvester\", \"hatchet\", \"holster\", \"home theater\", \"honeycomb\", \"hook\", \"hoop skirt\", \"gymnastic horizontal bar\", \"horse-drawn vehicle\", \"hourglass\", \"iPod\", \"clothes iron\", \"carved pumpkin\", \"jeans\", \"jeep\", \"T-shirt\", \"jigsaw puzzle\", \"rickshaw\", \"joystick\", \"kimono\", \"knee pad\", \"knot\", \"lab coat\", \"ladle\", \"lampshade\", \"laptop computer\", \"lawn mower\", \"lens cap\", \"letter opener\", \"library\", \"lifeboat\", \"lighter\", \"limousine\", \"ocean liner\", \"lipstick\", \"slip-on shoe\", \"lotion\", \"music speaker\", \"loupe magnifying glass\", \"sawmill\", \"magnetic compass\", \"messenger bag\", \"mailbox\", \"tights\", \"one-piece bathing suit\", \"manhole cover\", \"maraca\", \"marimba\", \"mask\", \"matchstick\", \"maypole\", \"maze\", \"measuring cup\", \"medicine cabinet\", \"megalith\", \"microphone\", \"microwave oven\", \"military uniform\", \"milk can\", \"minibus\", \"miniskirt\", \"minivan\", \"missile\", \"mitten\", \"mixing bowl\", \"mobile home\", \"ford model t\", \"modem\", \"monastery\", \"monitor\", \"moped\", \"mortar and pestle\", \"graduation cap\", \"mosque\", \"mosquito net\", \"vespa\", \"mountain bike\", \"tent\", \"computer mouse\", \"mousetrap\", \"moving van\", \"muzzle\", \"metal nail\", \"neck brace\", \"necklace\", \"baby pacifier\", \"notebook computer\", \"obelisk\", \"oboe\", \"ocarina\", \"odometer\", \"oil filter\", \"pipe organ\", \"oscilloscope\", \"overskirt\", \"bullock cart\", \"oxygen mask\", \"product packet / packaging\", \"paddle\", \"paddle wheel\", \"padlock\", \"paintbrush\", \"pajamas\", \"palace\", \"pan flute\", \"paper towel\", \"parachute\", \"parallel bars\", \"park bench\", \"parking meter\", \"railroad car\", \"patio\", \"payphone\", \"pedestal\", \"pencil case\", \"pencil sharpener\", \"perfume\", \"Petri dish\", \"photocopier\", \"plectrum\", \"Pickelhaube\", \"picket fence\", \"pickup truck\", \"pier\", \"piggy bank\", \"pill bottle\", \"pillow\", \"ping-pong ball\", \"pinwheel\", \"pirate ship\", \"drink pitcher\", \"block plane\", \"planetarium\", \"plastic bag\", \"plate rack\", \"farm plow\", \"plunger\", \"Polaroid camera\", \"pole\", \"police van\", \"poncho\", \"pool table\", \"soda bottle\", \"plant pot\", \"potter's wheel\", \"power drill\", \"prayer rug\", \"printer\", \"prison\", \"missile\", \"projector\", \"hockey puck\", \"punching bag\", \"purse\", \"quill\", \"quilt\", \"race car\", \"racket\", \"radiator\", \"radio\", \"radio telescope\", \"rain barrel\", \"recreational vehicle\", \"fishing casting reel\", \"reflex camera\", \"refrigerator\", \"remote control\", \"restaurant\", \"revolver\", \"rifle\", \"rocking chair\", \"rotisserie\", \"eraser\", \"rugby ball\", \"ruler measuring stick\", \"sneaker\", \"safe\", \"safety pin\", \"salt shaker\", \"sandal\", \"sarong\", \"saxophone\", \"scabbard\", \"weighing scale\", \"school bus\", \"schooner\", \"scoreboard\", \"CRT monitor\", \"screw\", \"screwdriver\", \"seat belt\", \"sewing machine\", \"shield\", \"shoe store\", \"shoji screen / room divider\", \"shopping basket\", \"shopping cart\", \"shovel\", \"shower cap\", \"shower curtain\", \"ski\", \"balaclava ski mask\", \"sleeping bag\", \"slide rule\", \"sliding door\", \"slot machine\", \"snorkel\", \"snowmobile\", \"snowplow\", \"soap dispenser\", \"soccer ball\", \"sock\", \"solar thermal collector\", \"sombrero\", \"soup bowl\", \"keyboard space bar\", \"space heater\", \"space shuttle\", \"spatula\", \"motorboat\", \"spider web\", \"spindle\", \"sports car\", \"spotlight\", \"stage\", \"steam locomotive\", \"through arch bridge\", \"steel drum\", \"stethoscope\", \"scarf\", \"stone wall\", \"stopwatch\", \"stove\", \"strainer\", \"tram\", \"stretcher\", \"couch\", \"stupa\", \"submarine\", \"suit\", \"sundial\", \"sunglasses\", \"sunglasses\", \"sunscreen\", \"suspension bridge\", \"mop\", \"sweatshirt\", \"swim trunks / shorts\", \"swing\", \"electrical switch\", \"syringe\", \"table lamp\", \"tank\", \"tape player\", \"teapot\", \"teddy bear\", \"television\", \"tennis ball\", \"thatched roof\", \"front curtain\", \"thimble\", \"threshing machine\", \"throne\", \"tile roof\", \"toaster\", \"tobacco shop\", \"toilet seat\", \"torch\", \"totem pole\", \"tow truck\", \"toy store\", \"tractor\", \"semi-trailer truck\", \"tray\", \"trench coat\", \"tricycle\", \"trimaran\", \"tripod\", \"triumphal arch\", \"trolleybus\", \"trombone\", \"hot tub\", \"turnstile\", \"typewriter keyboard\", \"umbrella\", \"unicycle\", \"upright piano\", \"vacuum cleaner\", \"vase\", \"vaulted or arched ceiling\", \"velvet fabric\", \"vending machine\", \"vestment\", \"viaduct\", \"violin\", \"volleyball\", \"waffle iron\", \"wall clock\", \"wallet\", \"wardrobe\", \"military aircraft\", \"sink\", \"washing machine\", \"water bottle\", \"water jug\", \"water tower\", \"whiskey jug\", \"whistle\", \"hair wig\", \"window screen\", \"window shade\", \"Windsor tie\", \"wine bottle\", \"airplane wing\", \"wok\", \"wooden spoon\", \"wool\", \"split-rail fence\", \"shipwreck\", \"sailboat\", \"yurt\", \"website\", \"comic book\", \"crossword\", \"traffic or street sign\", \"traffic light\", \"dust jacket\", \"menu\", \"plate\", \"guacamole\", \"consomme\", \"hot pot\", \"trifle\", \"ice cream\", \"popsicle\", \"baguette\", \"bagel\", \"pretzel\", \"cheeseburger\", \"hot dog\", \"mashed potatoes\", \"cabbage\", \"broccoli\", \"cauliflower\", \"zucchini\", \"spaghetti squash\", \"acorn squash\", \"butternut squash\", \"cucumber\", \"artichoke\", \"bell pepper\", \"cardoon\", \"mushroom\", \"Granny Smith apple\", \"strawberry\", \"orange\", \"lemon\", \"fig\", \"pineapple\", \"banana\", \"jackfruit\", \"cherimoya (custard apple)\", \"pomegranate\", \"hay\", \"carbonara\", \"chocolate syrup\", \"dough\", \"meatloaf\", \"pizza\", \"pot pie\", \"burrito\", \"red wine\", \"espresso\", \"tea cup\", \"eggnog\", \"mountain\", \"bubble\", \"cliff\", \"coral reef\", \"geyser\", \"lakeshore\", \"promontory\", \"sandbar\", \"beach\", \"valley\", \"volcano\", \"baseball player\", \"bridegroom\", \"scuba diver\", \"rapeseed\", \"daisy\", \"yellow lady's slipper\", \"corn\", \"acorn\", \"rose hip\", \"horse chestnut seed\", \"coral fungus\", \"agaric\", \"gyromitra\", \"stinkhorn mushroom\", \"earth star fungus\", \"hen of the woods mushroom\", \"bolete\", \"corn cob\", \"toilet paper\"]\n",
+        "label_int = dict(\n",
+        "    zip(imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES,\n",
+        "        range(len(imagenet_class_names.CLIP_IMAGENET_CLASS_NAMES))))[label]"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "_CiyXD_6nQbu"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%capture --no-display\n",
+        "batch_size = 8\n",
+        "\n",
+        "target_labels = jnp.full((batch_size,), label_int, jnp.int32)\n",
+        "\n",
+        "if is_ar:\n",
+        "  ar_generation_config = dict(cfg_inference_weight=cfg_w, temp=temp)\n",
+        "  masked_generation_config = None\n",
+        "else:\n",
+        "  ar_generation_config = {}\n",
+        "  masked_generation_config = parallel_decode.MaskedGenerationConfig(\n",
+        "      cfg_inference_weight=cfg_w,\n",
+        "      choice_temperature = temp,\n",
+        "      num_steps = 16,\n",
+        "      ordering = \"maskgit\",\n",
+        "      schedule = \"cosine\",\n",
+        "  )\n",
+        "\n",
+        "# Sample from GIVT and decode\n",
+        "sampled_codes, _ = sample(\n",
+        "    target_labels, jax.random.PRNGKey(rng),\n",
+        "    tuple(ar_generation_config.items()), masked_generation_config)\n",
+        "\n",
+        "generated_images = vae_decode(sampled_codes)"
+      ],
+      "metadata": {
+        "id": "sCcGB0m1oQY1"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown Visualize images\n",
+        "ncols = 4\n",
+        "nrows = generated_images.shape[0] // ncols\n",
+        "fig, axes = plt.subplots(nrows, ncols, figsize=(4 * ncols, 4 * nrows))\n",
+        "\n",
+        "for idx, (ax, img) in enumerate(zip(axes.flat, generated_images)):\n",
+        "  ax.imshow(img * .5 + .5)\n",
+        "  if idx == 0:\n",
+        "    ax.set_title(f'Label: {label} ({label_int})', fontsize=10, ha='left', loc='left')\n",
+        "  ax.set_axis_off()"
+      ],
+      "metadata": {
+        "id": "4FWgfAghuh8P",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown Visualize latent codes\n",
+        "nrows = sampled_codes.shape[0]\n",
+        "ncols = sampled_codes.shape[-1] + 1\n",
+        "fig, axes = plt.subplots(nrows, ncols, figsize=(ncols, nrows))\n",
+        "\n",
+        "for r, (row_ax, code) in enumerate(zip(axes, sampled_codes)):\n",
+        "  code_norm = (code - code.min()) / (code.max() - code.min())\n",
+        "  for c, ax in enumerate(row_ax):\n",
+        "    if c == 0:\n",
+        "      cc = generated_images[r] * .5 + .5\n",
+        "    else:\n",
+        "      cc = code_norm[..., c - 1].reshape(res // 16, res // 16)\n",
+        "    ax.imshow(cc)\n",
+        "    ax.set_axis_off()"
+      ],
+      "metadata": {
+        "id": "zGPPeXONy0Am",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

big_vision_repo/big_vision/configs/proj/givt/givt_imagenet2012.py

@@ -0,0 +1,222 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Train Generative Infinite Vocabulary Transformer (GIVT) on ImageNet.
+
+Example launch command (local; see main README for launching on TPU servers):
+
+  python -m big_vision.trainers.proj.givt.generative \
+    --config big_vision/configs/proj/givt/givt_imagenet2012.py \
+    --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%m-%d_%H%M'`
+
+Add the suffix `:key1=value1,key2=value2,...` to the config path in the launch 
+command to modify the the config with the arguments below. For example:
+`--config big_vision/configs/proj/givt/givt_imagenet_2012.py:model_size=large`
+"""
+
+import big_vision.configs.common as bvcc
+import ml_collections
+
+
+RES = 256
+PATCH_SIZE = 16
+
+GIVT_MODELS = {
+    'base': dict(num_decoder_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768, dec_dropout_rate=0.1),
+    'default': dict(num_decoder_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024, dec_dropout_rate=0.2),
+    'large': dict(num_decoder_layers=48, num_heads=16, mlp_dim=8192, emb_dim=1536, dec_dropout_rate=0.3)
+}
+
+
+def get_config(arg=None):
+  """A config for training a simple VAE on imagenet2012."""
+  arg = bvcc.parse_arg(arg, res=RES, patch_size=PATCH_SIZE, style='ar',  # 'ar' or 'masked'
+                       model_size='default', runlocal=False, singlehost=False,
+                       adaptor=False)
+  config = ml_collections.ConfigDict()
+
+  config.input = {}
+  ### Using Imagenette here to ensure this config is runnable without manual
+  ### download of ImageNet. This is only meant for testing and will overfit
+  ### immediately. Please download ImageNet to reproduce the paper results.
+  # config.input.data = dict(name='imagenet2012', split='train[4096:]')
+  config.input.data = dict(name='imagenette', split='train')
+
+  config.input.batch_size = 8 * 1024 if not arg.runlocal else 8
+  config.input.shuffle_buffer_size = 25_000 if not arg.runlocal else 10
+
+  config.total_epochs = 500
+
+  config.input.pp = (
+      f'decode_jpeg_and_inception_crop({arg.res},'
+      f'area_min=80, area_max=100, ratio_min=1.0, ratio_max=1.0,'
+      f'method="bicubic", antialias=True)'
+      f'|flip_lr'
+      f'|value_range(-1, 1, key="image")'
+      f'|copy("label", "labels")'
+      f'|keep("image", "labels")')
+
+  pp_eval = (
+      f'decode'
+      f'|resize_small({arg.res}, inkey="image", outkey="image",'
+      f'method="bicubic", antialias=True)'
+      f'|central_crop({arg.res})'
+      f'|value_range(-1, 1, key="image")'
+      f'|copy("label", "labels")'
+      f'|keep("image", "labels")')
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None
+
+  # Flags for AR model.
+  config.ar_generation_config = ml_collections.ConfigDict()
+  config.ar_generation_config.temp = 0.95
+  config.ar_generation_config.temp_probs = 1.0
+  config.ar_generation_config.beam_size = 1
+  config.ar_generation_config.fan_size = 1
+  config.ar_generation_config.rand_top_k = False
+  config.ar_generation_config.rand_top_k_temp = 1.0
+  config.ar_generation_config.cfg_inference_weight = 0.4
+
+  # Flags for masked model.
+  config.masked_generation_config = ml_collections.ConfigDict()
+  config.masked_generation_config.choice_temperature = 35.0
+  config.masked_generation_config.ordering = 'maskgit'
+  config.masked_generation_config.cfg_inference_weight = 0.0
+  config.masked_generation_config.schedule = 'cosine'
+
+  # Used for eval sweep.
+  config.eval_only = False
+
+  # VAE section
+  config.vae = {}
+  config.vae.model = ml_collections.ConfigDict()
+  config.vae.model.code_len = (arg.res // arg.patch_size) ** 2
+  config.vae.model_name = 'proj.givt.cnn'
+  config.vae.model.codeword_dim = 16
+  config.vae.model.filters = 128
+  config.vae.model.num_res_blocks = 2
+  config.vae.model.channel_multipliers = (1, 1, 2, 2, 4)
+  config.vae.model.conv_downsample = False
+  config.vae.model.activation_fn = 'swish'
+  config.vae.model.norm_type = 'GN'
+  if arg.model_size == 'large':
+    config.vae.model_init = 'gs://big_vision/givt/vae_imagenet_2012_beta_1e-5_params'
+  else:
+    config.vae.model_init = 'gs://big_vision/givt/vae_imagenet_2012_beta_5e-5_params'
+  config.vae.model.malib_ckpt = True
+  config.vae.model_load = {}
+  config.vae.model_load.malib_ckpt = config.vae.model.malib_ckpt
+  config.vae.model_load.use_ema_params = True
+
+  # GIVT section
+  config.model_name = 'proj.givt.givt'
+  config.model_init = ''
+  assert arg.model_size in GIVT_MODELS, f'Unknown model size: {arg.model_size}'
+  config.model = ml_collections.ConfigDict(GIVT_MODELS[arg.model_size])
+  config.model.num_layers = 0
+  config.model.num_labels = 1000  # None
+  config.model.seq_len = config.vae.model.code_len
+  config.model.out_dim = config.vae.model.codeword_dim
+  config.model.num_mixtures = 16
+  config.model.posemb_type = 'learn'
+  config.model.scale_tol = 1e-6
+  config.model.style = arg.style
+  config.model.min_masking_rate_training = 0.3
+  config.model.mask_style = 'concat'
+  config.model.drop_labels_probability = 0.1
+  config.model.fix_square_plus = True
+  config.model.per_channel_mixtures = False
+  config.model_init = ''
+  # Required for model sharding
+  config.model.scan = True
+  config.model.remat_policy = 'nothing_saveable'
+
+  # Adaptor section
+  config.adaptor_name = 'proj.givt.adaptor' if arg.adaptor else ''
+  config.adaptor = {}
+  config.adaptor.model = ml_collections.ConfigDict()
+  config.adaptor.model.num_blocks = 8
+  config.adaptor.model.num_channels_bottleneck = 4 * config.model.out_dim
+
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(b2=0.95)
+  config.grad_clip_norm = 1.0
+
+  # FSDP training by default
+  config.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  config.sharding_rules = [('act_batch', ('data',))]
+
+  # Standard schedule
+  config.lr = 0.001
+  config.wd = 0.0001
+  config.schedule = dict(decay_type='cosine', warmup_percent=0.1)
+
+  # MaskGIT-specific parameters
+  if arg.style == 'masked':
+    config.model.dec_dropout_rate = 0.4
+    config.wd = 0.0
+    if arg.res == 512:
+      config.masked_generation_config.choice_temperature = 140
+  # GIVT-Causal 512px specific parameters
+  elif arg.res == 512 and arg.model_size == 'large':
+    config.model.dec_dropout_rate = 0.1
+    # Set up space-to-depth/pixel shuffle
+    config.vae.model.code_len //= 2
+    config.vae.model.pixel_shuffle_patch_size = (1, 2)
+    config.model.seq_len //= 2
+    config.model.out_dim = config.vae.model.codeword_dim * 2
+    config.model.num_mixtures = 32
+    config.adaptor.model.num_channels_bottleneck = 8 * config.model.out_dim
+    config.adaptor.model.pixel_shuffle_patch_size = (1, 2)
+    # Update sampling config
+    config.ar_generation_config.temp = 0.9
+    config.ar_generation_config.cfg_inference_weight = 0.9
+
+  ### Evaluation section
+  config.evals = {}
+  config.evals.val = ml_collections.ConfigDict()
+  config.evals.val.type = 'mean'
+  config.evals.val.pred = 'validation'
+  config.evals.val.data = {**config.input.data}
+  config.evals.val.data.split = f'train[:{4096 if not arg.runlocal else 8}]'
+  config.evals.val.pp_fn = pp_eval
+  config.evals.val.log_steps = 1_000 if not arg.runlocal else 20
+
+  config.evals.save_pred_sampling = dict(
+      type='proj.givt.save_predictions',
+      pp_fn=pp_eval,
+      log_steps=10_000,
+      pred='sample',
+      batch_size=512,
+      data=dict(name=config.input.data.name, split='validation[:512]'),
+      outfile='inference_sampled.npz',
+  )
+
+  config.seed = 0
+
+  config.ckpt_timeout = 30
+
+  if arg.runlocal:
+    config.input.batch_size = 4
+    config.input.shuffle_buffer_size = 10
+    config.log_training_steps = 5
+    config.model.num_decoder_layers = 2
+
+    config.evals.val.data.split = 'validation[:16]'
+    config.evals.val.log_steps = 20
+
+  return config

big_vision_repo/big_vision/configs/proj/givt/givt_nyu_depth.py

@@ -0,0 +1,198 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Train a GIVT encoder-decoder model for NYU depth prediction."""
+
+import itertools
+import big_vision.configs.common as bvcc
+import ml_collections
+
+ConfigDict = ml_collections.ConfigDict
+
+VTT_MODELS = {
+    'base': dict(num_layers=12, num_decoder_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768),
+    'large': dict(num_layers=24, num_decoder_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024),
+}
+
+RES = 512
+PATCH_SIZE = 16
+LABEL_RES = 512
+LABEL_PATCH_SIZE = 16
+QUANTIZATION_BINS = 256
+MIN_DEPTH = 0.001
+MAX_DEPTH = 10.0
+
+
+def get_config(arg='split=sweep'):
+  """Config for training."""
+  arg = bvcc.parse_arg(arg, split='sweep', runlocal=False, singlehost=False)
+  config = ConfigDict()
+
+  config.input = {}
+  config.input.pp = (
+      f'decode|nyu_depth|'
+      f'randu("fliplr")|det_fliplr(key="image")|det_fliplr(key="labels")|'
+      f'inception_box|crop_box(key="image")|crop_box(key="labels")|'
+      f'resize({RES})|'
+      f'resize({LABEL_RES},key="labels",method="nearest")|'
+      f'bin_nyu_depth(min_depth={MIN_DEPTH}, max_depth={MAX_DEPTH}, num_bins={QUANTIZATION_BINS})|'
+      f'value_range(-1,1)|'
+      f'copy("image", "cond_image")|copy("labels", "image")|'
+      f'keep("image", "cond_image")'
+  )
+  pp_eval = (
+      f'decode|nyu_depth|'
+      f'nyu_eval_crop|'
+      f'resize({RES})|'
+      f'resize({LABEL_RES},key="labels",method="nearest")|'
+      f'bin_nyu_depth(min_depth={MIN_DEPTH}, max_depth={MAX_DEPTH}, num_bins={QUANTIZATION_BINS})|'
+      f'value_range(-1,1)|'
+      f'copy("image", "cond_image")|copy("labels", "image")|'
+      f'keep("image", "cond_image")'
+  )
+  pp_predict = (
+      f'decode|nyu_depth|'
+      f'nyu_eval_crop|copy("labels","ground_truth")|'
+      f'resize({RES})|'
+      f'value_range(-1,1)|'
+      f'copy("image", "cond_image")|'
+      f'strong_hash(inkey="tfds_id", outkey="image/id")|'
+      f'keep("cond_image", "ground_truth", "image/id")'
+  )
+
+  config.input.data = dict(name='nyu_depth_v2', split='train')
+  config.input.batch_size = 512
+  config.input.shuffle_buffer_size = 50_000
+
+  config.total_epochs = 50
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None
+  config.prefetch_to_device = 2
+  config.seed = 0
+
+  # Optimizer section
+  config.optax_name = 'big_vision.scale_by_adafactor'
+  config.optax = dict(beta2_cap=0.95)
+
+  config.ar_generation_config = ConfigDict()
+  config.ar_generation_config.temp = 0.9
+  config.ar_generation_config.temp_probs = 1.0
+  config.ar_generation_config.beam_size = 2
+  config.ar_generation_config.fan_size = 8
+  config.ar_generation_config.rand_top_k = False
+  config.ar_generation_config.rand_top_k_temp = 1.0
+
+  config.lr = 0.001
+  config.wd = 0.000001
+  config.lr_mults = [
+      ('pos_embedding_encoder.*', 0.1),
+      ('EmbedPatches.*', 0.1),
+      ('encoder.*', 0.1),
+      ('decoder.*', 1.0)
+  ]
+  config.schedule = dict(decay_type='cosine', warmup_percent=0.1)
+
+  # Oracle section
+  config.min_depth = MIN_DEPTH
+  config.max_depth = MAX_DEPTH
+  config.vae = ConfigDict()
+  config.vae.model_name = 'proj.givt.vit'
+  config.vae.model = ConfigDict()
+  config.vae.model.input_size = (RES, RES)
+  config.vae.model.patch_size = (PATCH_SIZE, PATCH_SIZE)
+  config.vae.model.code_len = 256
+  config.vae.model.width = 768
+  config.vae.model.enc_depth = 6
+  config.vae.model.dec_depth = 12
+  config.vae.model.mlp_dim = 3072
+  config.vae.model.num_heads = 12
+  config.vae.model.codeword_dim = 16
+  config.vae.model.code_dropout = 'none'
+  config.vae.model.bottleneck_resize = True
+  # values: (channel index in source image, number of classes)
+  config.vae.model.inout_specs = {
+      'depth': (0, QUANTIZATION_BINS),
+  }
+  config.vae.model_init = 'gs://big_vision/givt/vae_nyu_depth_params.npz'
+
+  # Model section
+  config.model_name = 'proj.givt.givt'
+  # # Base model (for exploration)
+  # config.model_init = {'encoder': 'howto-i21k-B/16'}
+  # config.model = ConfigDict(VTT_MODELS['base'])
+  # Large model
+  config.model_init = {'encoder': 'howto-i21k-L/16'}
+  config.model_load = dict(dont_load=('cls', 'head/bias', 'head/kernel'))
+  config.model = ConfigDict(VTT_MODELS['large'])
+  config.model.patches = (PATCH_SIZE, PATCH_SIZE)
+  config.model.input_size = (RES, RES)
+  config.model.posemb_type = 'learn'
+  config.model.seq_len = config.vae.model.code_len
+  config.model.num_labels = None
+  config.model.num_mixtures = 1
+  config.model.fix_square_plus = True
+  config.model.out_dim = config.vae.model.codeword_dim
+  config.model.scale_tol = 1e-6
+  config.model.dec_dropout_rate = 0.0
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = ConfigDict()
+  config.evals.val.type = 'mean'
+  config.evals.val.pred = 'validation'
+  config.evals.val.data = {**config.input.data}
+  config.evals.val.data.split = 'validation'
+  config.evals.val.pp_fn = pp_eval
+  config.evals.val.log_steps = 250
+
+  base = {
+      'type': 'proj.givt.nyu_depth',
+      'data': {**config.input.data},
+      'pp_fn': pp_predict,
+      'pred': 'sample_depth',
+      'log_steps': 2000,
+      'min_depth': MIN_DEPTH,
+      'max_depth': MAX_DEPTH,
+  }
+
+  config.evals.nyu_depth_val = dict(base)
+  config.evals.nyu_depth_val.data.split = 'validation'
+
+  config.evals.save_pred = dict(base)
+  config.evals.save_pred.type = 'proj.givt.save_predictions'
+  del config.evals.save_pred.min_depth, config.evals.save_pred.max_depth
+  config.evals.save_pred.log_steps = 100_000
+  config.evals.save_pred.data.split = 'validation[:128]'
+  config.evals.save_pred.outfile = 'inference.npz'
+
+  config.eval_only = False
+  config.seed = 0
+
+  if arg.runlocal:
+    config.input.batch_size = 4
+    config.input.shuffle_buffer_size = 10
+    config.evals.val.log_steps = 20
+    config.evals.val.data.split = 'validation[:4]'
+    config.evals.nyu_depth_val.data.split = 'validation[:4]'
+    config.evals.save_pred.data.split = 'validation[:4]'
+    config.model.update(VTT_MODELS['base'])
+    del config.model_init
+    for k in config.evals.keys():
+      if k not in ['val', 'nyu_depth_val', 'save_pred']:
+        del config.evals[k]
+
+  return config

big_vision_repo/big_vision/configs/proj/givt/vae_coco_panoptic.py

@@ -0,0 +1,136 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Train VAE for GIVT-based UViM COCO panoptic task.
+"""
+
+import big_vision.configs.common as bvcc
+import ml_collections as mlc
+
+
+def get_config(arg='res=512,patch_size=16'):
+  """Config for training label compression on COCO-panoptic."""
+  arg = bvcc.parse_arg(arg, res=512, patch_size=16,
+                       runlocal=False, singlehost=False)
+  config = mlc.ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(name='coco/2017_panoptic', split='train[4096:]')
+
+  config.input.batch_size = 1024
+  config.input.shuffle_buffer_size = 25_000
+
+  config.total_epochs = 500
+
+  config.input.pp = (
+      f'decode|coco_panoptic|concat(["semantics","instances"], "labels")|'
+      f'randu("fliplr")|det_fliplr(key="image")|det_fliplr(key="labels")|'
+      f'inception_box|crop_box(key="image")|crop_box(key="labels")|'
+      f'resize({arg.res})|resize({arg.res},key="labels",method="nearest")|'
+      f'value_range(-1, 1)|make_canonical|copy("labels","image")|keep("image")'
+  )
+  pp_eval = (
+      f'decode|coco_panoptic|concat(["semantics","instances"], "labels")|'
+      f'resize({arg.res})|resize({arg.res},key="labels",method="nearest")|'
+      f'value_range(-1, 1)|make_canonical|copy("labels","image")|keep("image", "image/id")'
+  )
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None
+
+  # Model section
+  config.model_name = 'proj.givt.vit'
+  config.model = mlc.ConfigDict()
+  config.model.input_size = (arg.res, arg.res)
+  config.model.patch_size = (arg.patch_size, arg.patch_size)
+  config.model.code_len = 256
+  config.model.width = 768
+  config.model.enc_depth = 6
+  config.model.dec_depth = 12
+  config.model.mlp_dim = 3072
+  config.model.num_heads = 12
+  config.model.codeword_dim = 32
+  config.model.code_dropout = 'none'
+  config.model.bottleneck_resize = True
+  config.model.scan = True
+  config.model.remat_policy = 'nothing_saveable'
+
+  config.rec_loss_fn = 'xent'  # xent, l2
+  # values: (index in source image, number of classes)
+  config.model.inout_specs = {
+      'semantics': (0, 133 + 1),  # +1 for void label
+      'instances': (1, 100),  # COCO: actually 98 train/78 validation.
+  }
+
+  config.beta = 2.5e-4
+  config.beta_percept = 0.0
+
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(b2=0.95)
+  config.grad_clip_norm = 1.0
+
+  # FSDP training by default
+  config.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  config.sharding_rules = [('act_batch', ('data',))]
+
+  config.lr = 1e-3
+  config.wd = 1e-4
+  config.schedule = dict(decay_type='cosine', warmup_steps=0.1)
+  config.grad_clip_norm = 1.0
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = mlc.ConfigDict()
+  config.evals.val.type = 'mean'
+  config.evals.val.pred = 'validation'
+  config.evals.val.data = {**config.input.data}
+  config.evals.val.data.split = 'train[:4096]'
+  config.evals.val.pp_fn = pp_eval
+  config.evals.val.log_steps = 250
+
+  base = {
+      'type': 'proj.givt.coco_panoptic',
+      'pp_fn': pp_eval,
+      'log_steps': 5_000,
+      'pred': 'predict_panoptic',
+      # Filters objects that occupy less than 0.03^2 fraction of all pixels.
+      # 'pred_kw': {'min_fraction': 0.03 ** 2},
+  }
+  config.evals.coco_panoptic_train = dict(**base, data={'split': 'train[4096:8192]'})
+  config.evals.coco_panoptic_holdout = dict(**base, data={'split': 'train[:4096]'})
+  config.evals.coco_panoptic = dict(**base, data={'split': 'validation'})
+
+  config.evals.save_pred = dict(type='proj.givt.save_predictions')
+  config.evals.save_pred.pp_fn = pp_eval
+  config.evals.save_pred.log_steps = 100_000
+  config.evals.save_pred.pred = 'predict_panoptic'
+  config.evals.save_pred.data = {**config.input.data}
+  config.evals.save_pred.data.split = 'validation[:1024]'
+  config.evals.save_pred.outfile = 'inference.npz'
+
+  config.seed = 0
+
+  if arg.singlehost:
+    config.input.batch_size = 128
+    config.num_epochs = 100
+  elif arg.runlocal:
+    config.input.batch_size = 16
+    config.input.shuffle_buffer_size = 10
+    config.log_training_steps = 5
+    config.model.enc_depth = 1
+    config.model.dec_depth = 1
+
+  return config

big_vision_repo/big_vision/configs/proj/givt/vae_nyu_depth.py

@@ -0,0 +1,158 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Train VAE on NYU depth data for GIVT-based UViM.
+"""
+
+import big_vision.configs.common as bvcc
+import ml_collections as mlc
+
+
+QUANTIZATION_BINS = 256
+MIN_DEPTH = 0.001
+MAX_DEPTH = 10.0
+
+
+def get_config(arg='res=512,patch_size=16'):
+  """Config for training label compression on NYU depth."""
+  arg = bvcc.parse_arg(arg, res=512, patch_size=16,
+                       runlocal=False, singlehost=False)
+  config = mlc.ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(name='nyu_depth_v2', split='train')
+
+  config.input.batch_size = 1024
+  config.input.shuffle_buffer_size = 25_000
+
+  config.total_epochs = 200
+
+  config.input.pp = (
+      f'decode|nyu_depth|'
+      f'randu("fliplr")|det_fliplr(key="image")|det_fliplr(key="labels")|'
+      f'inception_box|crop_box(key="image")|crop_box(key="labels")|'
+      f'resize({arg.res})|resize({arg.res},key="labels",method="nearest")|'
+      f'bin_nyu_depth(min_depth={MIN_DEPTH}, max_depth={MAX_DEPTH}, num_bins={QUANTIZATION_BINS})|'
+      f'value_range(-1, 1)|copy("labels", "image")|keep("image")'
+  )
+  pp_eval = (
+      f'decode|nyu_depth|nyu_eval_crop|'
+      f'resize({arg.res})|resize({arg.res},key="labels",method="nearest")|'
+      f'bin_nyu_depth(min_depth={MIN_DEPTH}, max_depth={MAX_DEPTH}, num_bins={QUANTIZATION_BINS})|'
+      f'value_range(-1, 1)|copy("labels", "image")|keep("image")'
+  )
+  pp_pred = (
+      f'decode|nyu_depth|nyu_eval_crop|copy("labels","ground_truth")|'
+      f'resize({arg.res})|resize({arg.res},key="labels",method="nearest")|'
+      f'bin_nyu_depth(min_depth={MIN_DEPTH}, max_depth={MAX_DEPTH}, num_bins={QUANTIZATION_BINS})|'
+      f'value_range(-1, 1)|copy("labels", "image")|'
+      f'keep("image", "ground_truth")'
+  )
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = None
+
+  # Model section
+  config.min_depth = MIN_DEPTH
+  config.max_depth = MAX_DEPTH
+  config.model_name = 'proj.givt.vit'
+  config.model = mlc.ConfigDict()
+  config.model.input_size = (arg.res, arg.res)
+  config.model.patch_size = (arg.patch_size, arg.patch_size)
+  config.model.code_len = 256
+  config.model.width = 768
+  config.model.enc_depth = 6
+  config.model.dec_depth = 12
+  config.model.mlp_dim = 3072
+  config.model.num_heads = 12
+  config.model.codeword_dim = 16
+  config.model.code_dropout = 'none'
+  config.model.bottleneck_resize = True
+  config.model.scan = True
+  config.model.remat_policy = 'nothing_saveable'
+  config.model_init = ''
+
+  config.rec_loss_fn = 'xent'  # xent, l2
+  config.mask_zero_target = True
+  # values: (index in source image, number of classes)
+  config.model.inout_specs = {
+      'depth': (0, QUANTIZATION_BINS),
+  }
+
+  config.beta = 2e-4
+  config.beta_percept = 0.0
+
+  # Optimizer section
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(b2=0.95)
+
+  # FSDP training by default
+  config.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  config.sharding_rules = [('act_batch', ('data',))]
+
+  config.lr = 1e-3
+  config.wd = 1e-4
+  config.schedule = dict(decay_type='cosine', warmup_steps=0.1)
+  config.grad_clip_norm = 1.0
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = mlc.ConfigDict()
+  config.evals.val.type = 'mean'
+  config.evals.val.pred = 'validation'
+  config.evals.val.data = {**config.input.data}
+  config.evals.val.data.split = 'validation'
+  config.evals.val.pp_fn = pp_eval
+  config.evals.val.log_steps = 250
+
+  base = {
+      'type': 'proj.givt.nyu_depth',
+      'data': {**config.input.data},
+      'pp_fn': pp_pred,
+      'pred': 'predict_depth',
+      'log_steps': 2000,
+      'min_depth': MIN_DEPTH,
+      'max_depth': MAX_DEPTH,
+  }
+  config.evals.nyu_depth_val = {**base}
+  config.evals.nyu_depth_val.data.split = 'validation'
+
+  # ### Uses a lot of memory
+  # config.evals.save_pred = dict(type='proj.givt.save_predictions')
+  # config.evals.save_pred.pp_fn = pp_eval
+  # config.evals.save_pred.log_steps = 100_000
+  # config.evals.save_pred.data = {**config.input.data}
+  # config.evals.save_pred.data.split = 'validation[:64]'
+  # config.evals.save_pred.batch_size = 64
+  # config.evals.save_pred.outfile = 'inference.npz'
+
+  config.eval_only = False
+  config.seed = 0
+
+  if arg.singlehost:
+    config.input.batch_size = 128
+    config.num_epochs = 50
+  elif arg.runlocal:
+    config.input.batch_size = 16
+    config.input.shuffle_buffer_size = 10
+    config.log_training_steps = 5
+    config.model.enc_depth = 1
+    config.model.dec_depth = 1
+    config.evals.val.data.split = 'validation[:16]'
+    config.evals.val.log_steps = 20
+    config.evals.nyu_depth_val.data.split = 'validation[:16]'
+
+  return config

big_vision_repo/big_vision/configs/proj/gsam/vit_i1k_gsam_no_aug.py

@@ -0,0 +1,134 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Pre-training ViT on ILSVRC-2012 with GSAM in https://arxiv.org/abs/2203.08065
+
+Run training of a B/32 model:
+
+big_vision.trainers.proj.gsam.train \
+    --config big_vision/configs/proj/gsam/vit_i1k_gsam_no_aug.py \
+    --workdir gs://[your_bucket]/big_vision/`date '+%m-%d_%H%M'`
+
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.common_fewshot import get_fewshot_lsr
+import ml_collections as mlc
+
+def get_config(arg=None):
+  """Config for training."""
+  arg = bvcc.parse_arg(arg, variant='B/32', runlocal=False)
+  config = mlc.ConfigDict()
+
+  config.dataset = 'imagenet2012'
+  config.train_split = 'train[:99%]'
+  config.cache_raw = not arg.runlocal  # Needs up to 120GB of RAM!
+  config.shuffle_buffer_size = 250_000  # Per host, so small-ish is ok.
+  config.num_classes = 1000
+  config.loss = 'sigmoid_xent'
+  config.batch_size = 4096
+  config.num_epochs = 300
+
+  pp_common = (
+      '|value_range(-1, 1)'
+      '|onehot(1000, key="{lbl}", key_result="labels")'
+      '|keep("image", "labels")'
+  )
+  config.pp_train = (
+      'decode_jpeg_and_inception_crop(224)|flip_lr|' +
+      pp_common.format(lbl='label')
+  )
+  pp = 'decode|resize_small(256)|central_crop(224)' + pp_common
+
+  # Aggressive pre-fetching because our models here are small, so we not only
+  # can afford it, but we also need it for the smallest models to not be
+  # bottle-necked by the input pipeline. Play around with it for -L models tho.
+  config.prefetch_to_host = 8
+  config.prefetch_to_device = 4
+
+  config.log_training_steps = 50
+  config.checkpoint_steps = 1000
+
+  # Model section
+  config.model_name = 'vit'
+  config.model = dict(
+      variant=arg.variant,
+      rep_size=False,
+      pool_type='gap',
+  )
+  config.init_head_bias = -10.0
+
+  # Optimizer section
+  config.grad_clip_norm = 1.0
+  config.optax_name = 'scale_by_adam'
+  config.optax = dict(mu_dtype='float32')
+  # The modified AdaFactor we introduced in https://arxiv.org/abs/2106.04560
+  # almost always behaves exactly like adam, but at a fraction of the memory
+  # cost (specifically, adam_bf16 = +1.5M, adafactor = +0.5M), hence it is a
+  # good idea to try it when you are memory-bound!
+  # config.optax_name = 'big_vision.scale_by_adafactor'
+  # A good flag to play with when hitting instabilities, is the following:
+  # config.optax = dict(beta2_cap=0.95)
+
+  config.lr = 0.003
+  config.wd = 0.001 # default is 0.0001; paper used 0.3, effective wd=0.3*lr
+  config.schedule = dict(
+      warmup_steps=10_000,
+      decay_type='linear',
+      linear_end=0.01,
+  )
+
+  # GSAM settings.
+  # Note: when rho_max=rho_min and alpha=0, GSAM reduces to SAM.
+  config.gsam = dict(
+      rho_max=0.6,
+      rho_min=0.1,
+      alpha=0.6,
+      lr_max=config.get_ref('lr'),
+      lr_min=config.schedule.get_ref('linear_end') * config.get_ref('lr'),
+  )
+
+  # Eval section
+  eval_common = dict(
+      type='classification',
+      dataset='imagenet2012',
+      pp_fn=pp.format(lbl='label'),
+      loss_name=config.loss,
+      log_steps=2500,  # Very fast O(seconds) so it's fine to run it often.
+  )
+  config.evals = {}
+  config.evals.train = {**eval_common, 'split': 'train[:2%]'}
+  config.evals.minival = {**eval_common, 'split': 'train[99%:]'}
+  config.evals.val = {**eval_common, 'split': 'validation'}
+  config.evals.v2 = {**eval_common, 'dataset': 'imagenet_v2', 'split': 'test'}
+
+  config.evals.real = {**eval_common}
+  config.evals.real.dataset = 'imagenet2012_real'
+  config.evals.real.split = 'validation'
+  config.evals.real.pp_fn = pp.format(lbl='real_label')
+
+  config.fewshot = get_fewshot_lsr(runlocal=arg.runlocal)
+  config.fewshot.log_steps = 10_000
+
+  # Make a few things much smaller for quick local debugging testruns.
+  if arg.runlocal:
+    config.shuffle_buffer_size = 10
+    config.batch_size = 8
+    config.minival.split = 'train[:16]'
+    config.val.split = 'validation[:16]'
+    config.real.split = 'validation[:16]'
+    config.v2.split = 'test[:16]'
+
+  return config

big_vision_repo/big_vision/configs/proj/image_text/README.md

@@ -0,0 +1,65 @@
+# Image/text models
+
+## LiT: Zero-Shot Transfer with Locked-image text Tuning
+
+*by Xiaohua Zhai, Xiao Wang, Basil Mustafa, Andreas Steiner, Daniel Keysers, Alexander Kolesnikov, Lucas Beyer*
+
+https://arxiv.org/abs/2111.07991
+
+```
+@article{zhai2022lit,
+  title={LiT: Zero-Shot Transfer with Locked-image Text Tuning},
+  author={Zhai, Xiaohua and Wang, Xiao and Mustafa, Basil and Steiner, Andreas and Keysers, Daniel and Kolesnikov, Alexander and Beyer, Lucas},
+  journal={CVPR},
+  year={2022}
+}
+```
+
+Model card:
+https://github.com/google-research/vision_transformer/blob/main/model_cards/lit.md
+
+Colabs:
+
+- https://colab.research.google.com/github/google-research/vision_transformer/blob/main/lit.ipynb
+- https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/lit.ipynb
+
+### Results
+
+| Model | Download link | ImageNet 0-shot | MS-COCO I→T | MS-COCO T→I | Config `arg` |
+| :---  | :---:         | :---:           | :---:       | :---:       | :---         |
+| mixed_L16L | [link](https://storage.googleapis.com/vit_models/lit/LiT-L16L.npz) | 75.7 | 48.5 | 31.2 | `txt=bert_large,img=L/16` |
+| mixed_B16B | [link](https://storage.googleapis.com/vit_models/lit/LiT-B16B.npz) | 72.1 | 49.4 | 31.1 | `txt=bert_base,img=B/16,img_head` |
+| mixed_B16B_2 | [link](https://storage.googleapis.com/vit_models/lit/LiT-B16B.npz) | 73.9 | 51.5 | 31.8 | `txt=bert_base,img=B/16` |
+| coco_B16B | [link](https://storage.googleapis.com/vit_models/lit/big_vision/coco_B16B/checkpoint.npz) | 20.7 | 47.2 | 32.1 | `txt=bert_base,img=B/16` |
+
+The first three rows are the best available models trained on open source data,
+originally published in the [`google-research/vision_transformer`] repository.
+These models were re-evaluated with this codebase using the following commands:
+
+```bash
+big_vision.tools.eval_only --config big_vision/configs/proj/image_text/lit_coco.py:txt=bert_base,img=B/16,img_head,init=gs://vit_models/lit/LiT-B16B.npz
+
+big_vision.tools.eval_only --config big_vision/configs/proj/image_text/lit_coco.py:txt=bert_base,img=B/16_2,init=gs://vit_models/lit/LiT-B16B_2.npz
+
+big_vision.tools.eval_only --config big_vision/configs/proj/image_text/lit_coco.py:txt=bert_large,img=L/16,init=gs://vit_models/lit/LiT-L16L.npz
+```
+
+Unfortunately, the public multi-modal datasets [`CC12M`] and [`YFCC100M`] are
+not yet available in [`tfds`], so these models cannot be reproduced with the
+codebase. For this reason we provide the much weaker model `coco_B16B` in the
+third row, which was trained on the small `tfds` dataset [`coco_captions`], and
+can be used to verify correctness of the codebase
+([workdir](https://console.cloud.google.com/storage/browser/vit_models/lit/big_vision/coco_B16B/)).
+
+[`google-research/vision_transformer`]: https://github.com/google-research/vision_transformer
+[`CC12M`]: https://arxiv.org/abs/2102.08981
+[`YFCC100M`]: https://arxiv.org/abs/1503.01817
+[`tfds`]: https://www.tensorflow.org/datasets/api_docs/python/tfds
+[`coco_captions`]: https://www.tensorflow.org/datasets/catalog/coco_captions
+
+
+### Changelog
+
+- 2022-08-18: Added LiT-B16B_2 model that was trained for 60k steps
+  (LiT_B16B: 30k) without linear head on the image side (LiT_B16B: 768) and has
+  better performance.

big_vision_repo/big_vision/configs/proj/image_text/common.py

@@ -0,0 +1,127 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Snippets and constants used a lot in image-text configs."""
+
+import ml_collections
+
+
+# pylint: disable=line-too-long
+inits = {
+    # Downloaded & extracted from original repo:
+    # https://github.com/google-research/bert
+    'bert_base': ('base', 'gs://vit_models/lit/bert/uncased_L-12_H-768_A-12'),
+    'bert_large': ('large', 'gs://vit_models/lit/bert/uncased_L-uncased_L-24_H-1024_A-16'),
+    # Recommended "How to train your ViT..." checkpoints from
+    # https://github.com/google-research/vision_transformer#available-vit-models
+    'B/32': ('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'),
+    'B/16': ('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'),
+    'L/16': ('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'),
+}
+# pylint: enable=line-too-long
+
+
+def _square875(sz):
+  return f'resize({int(sz/0.875)})|central_crop({sz})|value_range(-1,1)'
+
+
+def _aspect75(sz):
+  return f'resize_small({int(sz/0.75)})|central_crop({sz})|value_range(-1,1)'
+
+
+def _drop_no_real_label(f):
+  return len(f['real_label']) > 0
+
+
+def _drop_no_imagenet(f):
+  return len(f['labels_imagenet']) > 0
+
+
+DISCLF_DATASET_OVERRIDES = {
+    'imagenet2012': {'class_names': 'clip', 'split': 'validation'},
+    'imagenet2012_minival': {
+        'dataset_name': 'imagenet2012',
+        'class_names': 'clip',
+        'split': 'train[99%:]',
+    },
+    'imagenet2012_real': {
+        'split': 'validation',
+        'class_names': 'clip',
+        'class_names_dataset_name': 'imagenet2012',
+        'pp_img': lambda sz: (
+            _square875(sz) + '|pad_to_shape(inkey="real_label", outkey="label", shape=[10], pad_value=-1)|keep("label", "image")'),  # pylint: disable=line-too-long
+        'pre_filter_fn': _drop_no_real_label,
+    },
+    'imagenet_v2': {'class_names': 'clip'},
+    'imagenet_a': {
+        'class_names': 'clip',
+        'pp_img': lambda sz: _aspect75(sz) + '|map("i1k_i1ka")',
+    },
+    'imagenet_r': {
+        'class_names': 'clip',
+        'pp_img': lambda sz: _square875(sz) + '|map("i1k_i1kr")',
+    },
+}
+
+
+def get_disclf(sz, *, pp_txt=None, dataset_names=('imagenet2012',), **kw):
+  """Returns config for discriminative_classifier of specified datasets."""
+  config = ml_collections.ConfigDict(dict(
+      dataset_names=list(dataset_names),
+      type='proj.image_text.discriminative_classifier',
+      prefix='z/0shot/',
+      pp_img=_square875(sz),
+      dataset_overrides={},
+      cache_final=True,
+      **kw,
+  ))
+  if pp_txt:
+    config.pp_txt = pp_txt
+  for name in dataset_names:
+    if name in DISCLF_DATASET_OVERRIDES:
+      config.dataset_overrides[name] = {**DISCLF_DATASET_OVERRIDES[name]}
+      d = config.dataset_overrides[name]
+      if 'pp_img' in d and callable(d['pp_img']):
+        with d.ignore_type():
+          d['pp_img'] = d['pp_img'](sz)
+  return config
+
+
+def get_coco(
+    *,
+    pp_img='resize(224)|value_range(-1, 1)',
+    pp_txt='tokenize(max_len=16, inkey="texts", eos="sticky", pad_value=1)',
+    prefix='z/retr/coco_',
+    **kw):
+  """Returns config for mscoco retrieval zero-shot.
+
+  Args:
+    pp_img: Pre-processing string for "image" feature.
+    pp_txt: Pre-processing string for texts (expected to tokenize "texts" to
+      "labels").
+    prefix: Prefix to use for metrics.
+    **kw: Other config settings, most notably log_{steps,percent,...}.
+
+  Returns:
+    `ConfigDict` that can be used as a retrieval evaluator configuration.
+  """
+  return ml_collections.ConfigDict({
+      'type': 'proj.image_text.retrieval',
+      'pp_txt': pp_txt,
+      'pp_img': pp_img,
+      'prefix': prefix,
+      'dataset': 'coco_captions',
+      'txt_name': ('captions', 'text'),
+      **kw,
+  })

big_vision_repo/big_vision/configs/proj/image_text/siglip_lit_coco.py

@@ -0,0 +1,115 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""Minimal SigLIP (https://arxiv.org/abs/2303.15343) example.
+
+Example training:
+
+big_vision.trainers.proj.image_text.siglip \
+    --config big_vision/configs/proj/image_text/lit_coco.py:batch_size=512 \
+    --workdir gs://$GS_BUCKET_NAME/big_vision/`date '+%Y-%m-%d_%H%M'`
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.image_text import common
+from ml_collections import ConfigDict
+
+
+def get_config(arg=None):
+  """The base configuration."""
+  arg = bvcc.parse_arg(
+      arg, res=224, runlocal=False, token_len=16, txt='bert_base', img='B/16',
+      init='', img_head=False, batch_size=512)
+  img_name, img_init = common.inits[arg.img]
+  txt_name, txt_init = common.inits[arg.txt]
+  config = ConfigDict()
+
+  config.input = {}
+  config.input.data = dict(name='coco_captions', split='train')
+  config.input.batch_size = arg.batch_size if not arg.runlocal else 32
+  config.input.shuffle_buffer_size = 250_000  if not arg.runlocal else 50
+
+  config.total_steps = 5_000 if not arg.runlocal else 1
+
+  config.init_shapes = [(1, arg.res, arg.res, 3), (1, arg.token_len,)]
+  config.init_types = ['float32', 'int32']
+
+  if arg.init:
+    vocab_path = arg.init.rsplit('.', 1)[0] + '.txt'
+  else:
+    vocab_path = f'{txt_init}/vocab.txt'
+  tokenizer = lambda inkey: (
+      f'bert_tokenize(inkey="{inkey}", max_len={arg.token_len}, '
+      f'vocab_path="{vocab_path}")')
+  config.input.pp = (
+      f'decode|resize({arg.res})|flip_lr|randaug(2,10)|value_range(-1,1)'
+      f'|flatten|{tokenizer("captions/text")}|keep("image", "labels")'
+  )
+  config.pp_modules = ['ops_general', 'ops_image', 'ops_text',
+                       'proj.flaxformer.bert_ops', 'archive.randaug']
+
+  config.log_training_steps = 50
+  config.ckpt_steps = 1000
+
+  # Model section
+  config.model_name = 'proj.image_text.two_towers'
+  config.model_load = {}
+  if arg.init:
+    config.model_init = arg.init
+  else:
+    config.model_init = {'image': img_init, 'text': txt_init}
+    config.model_load['txt_load_kw'] = {'dont_load': ['head/kernel', 'head/bias']}
+    if not arg.img_head:
+      config.model_load['img_load_kw'] = {'dont_load': ['head/kernel', 'head/bias']}
+  config.model = ConfigDict()
+  config.model.image_model = 'vit'
+  config.model.text_model = 'proj.flaxformer.bert'
+  config.model.image = ConfigDict({
+      'variant': img_name,
+      'pool_type': 'tok',
+      'head_zeroinit': False,
+  })
+  config.model.text = ConfigDict({
+      'config': txt_name,
+      'head_zeroinit': False,
+  })
+  config.model.temperature_init = 10.0
+  dim = {'B': 768, 'L': 1024}[arg.img[0]]
+  config.model.out_dim = (dim if arg.img_head else None, dim)  # (image_out_dim, text_out_dim)
+  config.model.bias_init = -2.71
+
+  if txt_name == 'base':
+    config.optax_name = 'scale_by_adam'
+  else:
+    config.optax_name = 'big_vision.scale_by_adafactor'
+
+  config.lr = 0.001
+  config.wd = 0.01
+  warmup_steps = max(int(0.03 * config.total_steps), 100)
+  config.schedule = [
+      ('img/.*', None),  # Freezes image tower.
+      ('.*', dict(decay_type='cosine', warmup_steps=warmup_steps)),
+  ]
+
+  config.grad_clip_norm = 1.0
+
+  config.evals = {}
+  config.evals.retrieval_coco = common.get_coco(
+      pp_img=f'resize({arg.res})|value_range(-1, 1)',
+      pp_txt=tokenizer('texts'),
+      log_steps=1000,
+  )
+
+  return config

big_vision_repo/big_vision/configs/proj/paligemma/README.md

@@ -0,0 +1,270 @@
+# PaliGemma model README
+
+PaliGemma is an open vision-language model (VLM) inspired by PaLI-3, built with
+open components, such as
+the [SigLIP vision model](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP_demo.ipynb)
+and
+the [Gemma language model](https://ai.google.dev/gemma).
+PaliGemma is designed as a versatile model for transfer to a wide range of
+vision-language tasks such as image and short video caption, visual question
+answering, text reading, object detection and object segmentation. Together with
+the pretrained and transfer checkpoints at multiple resolutions, we provide a
+checkpoint transferred to a mixture of tasks that can be used for off-the-shelf
+exploration.
+
+## Quick Reference
+
+This is the reference repository of the model, you may also want to check out the resources on
+
+ - [ArXiv](https://arxiv.org/abs/2407.07726): Technical report.
+ - [Kaggle](https://www.kaggle.com/models/google/paligemma):
+ All pre-trained / mix checkpoints and model card.
+ - [Kaggle-FT](https://www.kaggle.com/models/google/paligemma-ft):
+ All fine-tuned checkpoints and model card.
+ - [VertexAI Model Garden](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/363):
+ Paligemma models on GCP.
+ - [Hugging Face](https://huggingface.co/google/paligemma-3b-pt-224):
+ PyTorch port of paligemma models.
+ - [Light finetuning colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/paligemma/finetune_paligemma.ipynb):
+  Lightweight colab for fine-tuning PaliGemma. It can be run on a single T4 GPU (16GB)
+  available on free Colab.
+ - [HuggingFace demo](https://hf.co/spaces/google/paligemma): live demo.
+
+### Citation BibTeX
+
+```
+@article{beyer2024paligemma,
+      title={{PaliGemma: A versatile 3B VLM for transfer}},
+      author={Lucas Beyer and Andreas Steiner and André Susano Pinto and Alexander Kolesnikov and Xiao Wang and Daniel Salz and Maxim Neumann and Ibrahim Alabdulmohsin and Michael Tschannen and Emanuele Bugliarello and Thomas Unterthiner and Daniel Keysers and Skanda Koppula and Fangyu Liu and Adam Grycner and Alexey Gritsenko and Neil Houlsby and Manoj Kumar and Keran Rong and Julian Eisenschlos and Rishabh Kabra and Matthias Bauer and Matko Bošnjak and Xi Chen and Matthias Minderer and Paul Voigtlaender and Ioana Bica and Ivana Balazevic and Joan Puigcerver and Pinelopi Papalampidi and Olivier Henaff and Xi Xiong and Radu Soricut and Jeremiah Harmsen and Xiaohua Zhai},
+      year={2024},
+      journal={arXiv preprint arXiv:2407.07726}
+}
+```
+
+## Model description
+
+### Overview
+
+PaliGemma-3B is Vision-Language model that was inspired by the PaLI-3 recipe.
+It is built on SigLIP visual encoder (specifically, SigLIP-So400m/14) and the
+Gemma 2B language model. PaliGemma takes as input one or more images,
+which are turned into "soft tokens" by the SigLIP encoder, and input text
+(codenamed the "prefix") that is tokenized by Gemma's tokenizer. The image
+tokens and prefix tokens are concatenated (in this order) and passed to the
+Gemma decoder with full block-attention, which then generates an output text
+(the "suffix") auto-regressively with masked attention.
+
+![PaliGemma model](paligemma.png)
+
+### Training stages
+
+Similar to PaLI-3, PaliGemma's training consists of multiple stages:
+
+  - Stage 0: the unimodal pre-training. We use publicly available off-the-shelf
+  SigLIP and Gemma models which have been pre-trained unimodally by their
+  respective authors.
+  - Stage 1: multimodal pre-training. The combined PaliGemma model is now
+  pre-trained on a fully multimodal training dataset, this at a low resolution
+  of 224px² and prefix+suffix sequence length of 128 tokens. This results in
+  the first base model that we release.
+  - Stage 2: high-resolution pre-training. We continue pre-training of the
+  Stage 1 model at resolution 448px² with sequence length 512 tokens for a short
+  duration on the same multimodal training data, but re-weighted with more
+  emphasis on examples that make use of higher resolution or longer sequence
+  length. We repeat this once more at resolution 896px². This results in two
+  further "high res" base models that we also release.
+  - Stage 3: fine-tune. The base models are transferred to
+  specific tasks by fine-tuning. To facilitate further research and
+  reproducibility, we release checkpoints fine-tuned on most of the benchmarks
+  we evaluate on. We also provide a "mix" transfer model, fine-tuned on a wide
+  variety of data, for use in interactive demos.
+
+Most of the code examples, use-cases, and code release are about Stage 3:
+transferring to a task or dataset of interest to the user. 
+
+### Tokenizer
+
+PaliGemma uses the Gemma tokenizer with 256'000 tokens, but we further extend
+its vocabulary with 1024 entries that represent coordinates in normalized
+image-space (\<loc0000>...\<loc1023>), and another with 128 entries
+(\<seg000>...\<seg127>) that are codewords used by a lightweight
+referring-expression segmentation vector-quantized variational auto-encoder
+(VQ-VAE) with the architecture of [Ning et al. (2023)](https://arxiv.org/abs/2301.02229) and trained on OpenImages
+as in PaLI-3. While the `big_vision` codebase is flexible enough to extend
+tokenizers on-the-fly, we also provide a SentencePiece model file of the Gemma
+tokenizer with these additional tokens baked in, for the convenience of
+other codebases.
+
+## Checkpoints
+
+The PaliGemma models are released under the same open license as the Gemma
+models, and hence require manual acknowledgement of the license terms on kaggle:
+https://www.kaggle.com/models/google/paligemma. The reference checkpoints are
+available on
+[Kaggle](https://www.kaggle.com/models/google/paligemma),
+[VertexAI Model Garden](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/363) and
+[Hugging Face](https://huggingface.co/google/paligemma-3b-pt-224).
+
+### Pretrained checkpoints
+
+Use one of these checkpoints as initialization for fine-tuning:
+
+  - pt-224: Versatile pretrained model for tasks that do not require seeing
+  small details in the image.
+  Examples: natural image captioning and question-answering, detection and
+  segmentation of medium-large objects. This model was trained with
+  sequence length 128.
+  - pt-448: Versatile base model for mid/higher resolution tasks with access
+  to smaller details. Besides higher resolution, it has gotten more weight on
+  text reading, detection, and segmentation during its pre-training. Examples:
+  as above, plus detection, segmentation, text/diagram reading. This model was
+  trained with sequence length 512.
+  - pt-896: Further scaled-up version of pt-448, especially good at reading
+  very small texts as often found in documents and infographics. This model
+  was trained with sequence length 512.
+
+Besides the reference float32 checkpoint (11GB), we further provide
+bfloat16 and float16 variants of each, to reduce download and storage time.
+These are good for inference and frozen transfers, but full fine-tuning
+should happen in float32 or mixed precision.
+
+### Mixture checkpoint
+
+This checkpoint is trained on a mixture of all our transfer tasks,
+with a balancing intended to make it "nice to use" out of the box for
+predictions. This model is multilingual and should
+understand prompts in various languages, although English
+is still its "mother tongue".
+Questions can be asked in a natural way (including asking for a caption or
+reading the text), and detection and segmentation should still work with the
+structured `detect {things}` and `segment {things}` prompts as in the base model.
+
+  - mix-224: Similarly to pt-224, this model is good at many natural image
+  tasks that do not require high resolution. Unlike the raw pre-trained model,
+  however, it can be interacted with more freely. For example, ask it to
+  "describe this image in great detail, please" or "How many coins do you see
+  in the picture?". This model was trained with sequence length 256.
+  - mix-448: As above, but it is better at tasks that require higher-resolution
+  input. For example, one could ask it "what is written in the "sum" field?",
+  to "describe this figure", or to "what is the GDP of France?" when shown an
+  infographic of countries' GDPs. This model was trained with
+  sequence length 512.
+
+### Transfers results and checkpoints
+
+We provide checkpoints transferred to most of the tasks we evaluated
+transfer on, see the [kaggle page](https://www.kaggle.com/models/google/paligemma).
+These are intended for use when a specialised model corresponding
+to one of the tasks is needed, for academic research purposes only.
+Depending on the task, they may require a specialised preprocessing format.
+
+The transfer setup is reasonably unified, with the main factors of variation
+being the training duration, learning-rate, and whether or not to use dropout
+and label-smoothing. Details can be found in the corresponding config files or
+in an upcoming tech report.
+
+Importantly, none of these tasks or datasets are part of the pre-training data
+mixture, and their images are explicitly removed from the web-scale
+pretraining data.
+
+#### Captioning
+
+Benchmark (train split) | Metric (split) | pt-224 | pt-448 | pt-896
+-----------------------|----------------|--------|--------|--------
+[COCO captions](https://cocodataset.org/#home) (train+restval) | CIDEr (val) | 141.92 | 144.60 | 
+[NoCaps](https://nocaps.org/) (Eval of COCO captions transfer) | CIDEr (val) | 121.72 | 123.58 | 
+[COCO-35L](https://arxiv.org/abs/2205.12522) (train) | CIDEr dev (en / avg-34 / avg) | 139.2 / 115.8 / 116.4 | 141.2 / 118.0 / 118.6 | 
+[XM3600](https://arxiv.org/abs/2205.12522) (Eval of COCO-35L transfer) | CIDEr test (en / avg-35 / avg) | 78.1 / 41.3 / 42.4 | 80.0 / 41.9 / 42.9 | 
+[TextCaps](https://textvqa.org/textcaps/) (train) | CIDEr (val) | 127.48 | 153.94 | 
+[SciCap](https://arxiv.org/abs/2110.11624) (first sentence, no subfigure) (train+val) | CIDEr / BLEU-4 (test) | 162.25 / 0.192 | 181.49 / 0.211 | 
+[Screen2words](https://arxiv.org/abs/2108.03353) (train+dev) | CIDEr (test) | 117.57 | 119.59 | 
+[Widget Captioning](https://arxiv.org/abs/2010.04295) (train+dev) | CIDEr (test) | 136.07 | 148.36 | 
+
+#### Question Answering
+
+Benchmark (train split) | Metric (split) | pt-224 | pt-448 | pt-896
+-----------------------|----------------|--------|--------|--------
+[VQAv2](https://visualqa.org/index.html) (train+validation) | Accuracy (Test server - std) | 83.19 | 85.64 | 
+[MMVP](https://arxiv.org/abs/2401.06209) (Eval of VQAv2 transfer) | Paired Accuracy | 47.33 | 45.33 | 
+[POPE](https://arxiv.org/abs/2305.10355) (Eval of VQAv2 transfer) | Accuracy (random / popular / adversarial) | 87.80 / 85.87 / 84.27 | 88.23 / 86.77 / 85.90 | 
+[Objaverse Multiview](https://arxiv.org/abs/2311.17851) (Eval of VQAv2 transfer) | Cosine Similarity (USEv4) | 62.7 | 62.8 | 
+[OKVQA](https://okvqa.allenai.org/) (train) | Accuracy (val) | 63.54 | 63.15 | 
+[A-OKVQA](https://allenai.org/project/a-okvqa/home) (MC) (train+val) | Accuracy (Test server) | 76.37 | 76.90 | 
+[A-OKVQA](https://allenai.org/project/a-okvqa/home) (DA) (train+val) | Accuracy (Test server) | 61.85 | 63.22 | 
+[GQA](https://cs.stanford.edu/people/dorarad/gqa/about.html) (train_balanced+val_balanced) | Accuracy (testdev balanced) | 65.61 | 67.03 | 
+[xGQA](https://aclanthology.org/2022.findings-acl.196/) (Eval of GQA transfer) | Mean Accuracy (bn,de,en,id,ko,pt,ru,zh) | 58.37 | 59.07 | 
+[NLVR2](https://lil.nlp.cornell.edu/nlvr/) (train+dev) | Accuracy (test) | 90.02 | 88.93 | 
+[MaRVL](https://marvl-challenge.github.io/) (Eval of NLVR2 transfer) | Mean Accuracy (test) (id,sw,ta,tr,zh) | 80.57 | 76.78 | 
+[AI2D](https://allenai.org/data/diagrams) (train) | Accuracy (test) | 72.12 | 73.28 | 
+[ScienceQA](https://scienceqa.github.io/) (Img subset, no CoT) (train+val) | Accuracy (test) | 95.39 | 95.93 | 
+[RSVQA-LR](https://zenodo.org/records/6344334) (Non numeric) (train+val) | Mean Accuracy (test) | 92.65 | 93.11 | 
+[RSVQA-HR](https://zenodo.org/records/6344367) (Non numeric) (train+val) | Mean Accuracy (test/test2) | 92.61 / 90.58 | 92.79 / 90.54 | 
+[ChartQA](https://arxiv.org/abs/2203.10244) (human+aug)x(train+val) | Mean Relaxed Accuracy (test_human, test_aug) | 57.08 | 71.36 | 
+[VizWiz](https://vizwiz.org/tasks-and-datasets/vqa/) VQA (train+val) | Accuracy (Test server - std) | 73.7 | 75.52 | 
+[TallyQA](https://arxiv.org/abs/1810.12440) (train) | Accuracy (test_simple/test_complex) | 81.72 / 69.56 | 84.86 / 72.27 | 
+[OCR-VQA](https://ocr-vqa.github.io/) (train+val) | Accuracy (test) | 73.24 | 75.60 | 75.90
+[TextVQA](https://textvqa.org/) (train+val) | Accuracy (Test server - std) | 55.47 | 73.15 | 76.48
+[DocVQA](https://www.docvqa.org/) (train+val) | ANLS (Test server) | 43.74 | 78.02 | 84.77
+[Infographic VQA](https://openaccess.thecvf.com/content/WACV2022/papers/Mathew_InfographicVQA_WACV_2022_paper.pdf) (train+val) | ANLS (Test server) | 28.46 | 40.47 | 47.75
+[SceneText VQA](https://arxiv.org/abs/1905.13648) (train+val) | ANLS (Test server) | 63.29 | 81.82 | 84.40
+
+#### Segmentation
+
+Benchmark (train split) | Metric (split) | pt-224 | pt-448 | pt-896
+-----------------------|----------------|--------|--------|--------
+[RefCOCO](https://arxiv.org/abs/1608.00272) (combined refcoco, refcoco+, refcocog excluding val and test images) | MIoU (validation) refcoco / refcoco+ / refcocog | 73.40 / 68.32 / 67.65 | 75.57 / 69.76 / 70.17 | 76.94 / 72.18 / 72.22
+
+#### Video tasks (Caption/QA)
+
+Benchmark (train split) | Metric (split) | pt-224 | pt-448 | pt-896
+-----------------------|----------------|--------|--------|--------
+[MSR-VTT](https://www.microsoft.com/en-us/research/publication/msr-vtt-a-large-video-description-dataset-for-bridging-video-and-language/) (Captioning) | CIDEr (test) | 70.54 | 
+[MSR-VTT](https://www.microsoft.com/en-us/research/publication/msr-vtt-a-large-video-description-dataset-for-bridging-video-and-language/) (QA) | Accuracy (test) | 50.09 | 
+[ActivityNet](http://activity-net.org/) (Captioning)] | CIDEr (test) | 34.62 | 
+[ActivityNet](http://activity-net.org/) (QA) | Accuracy (test) | 50.78 | 
+[VATEX](https://eric-xw.github.io/vatex-website/about.html) (Captioning) | CIDEr (test) | 79.73 | 
+[MSVD](https://www.cs.utexas.edu/users/ml/clamp/videoDescription/) (QA) | Accuracy (test) | 60.22 |
+
+#### Mix model (finetune on mixture of transfer tasks)
+
+Benchmark | Metric (split) | mix-224 | mix-448
+----------|----------------|---------|---------
+[MMVP](https://arxiv.org/abs/2401.06209) | Paired Accuracy | 46.00 | 45.33
+[POPE](https://arxiv.org/abs/2305.10355) | Accuracy (random / popular / adversarial) | 88.00 / 86.63 / 85.67 | 89.37 / 88.40 / 87.47
+
+
+## How to run PaliGemma fine-tuning
+
+To run PaliGemma fine-tuning, set up the `big_vision` repository by following the
+main README file. Here we provide PaliGemma-specific instructions.
+
+Checkpoints can be downloaded from Kaggle. You need to create an account and acknowledge checkpoint usage policy. You can then download any checkpoint:
+
+```
+export KAGGLE_USERNAME=
+export KAGGLE_KEY=
+
+# See https://www.kaggle.com/models/google/paligemma for a full list of models.
+export MODEL_NAME=paligemma-3b-pt-224
+export CKPT_FILE=paligemma-3b-pt-224.npz
+
+mkdir ckpts/
+cd ckpts/
+
+curl -L -u $KAGGLE_USERNAME:$KAGGLE_KEY\
+  -o pt_224.npz \
+  https://www.kaggle.com/api/v1/models/google/paligemma/jax/$MODEL_NAME/1/download/$CKPT_FILE
+```
+
+As an example, we provide the `forkme.py` config that is based on the easily-adjustable jsonl data source:
+
+```
+BV_GEMMA_DIR=ckpts/ python -m big_vision.trainers.proj.paligemma.train --config big_vision/configs/proj/paligemma/transfers/forkme.py --workdir workdirs/`date '+%m-%d_%H%M'`
+```
+
+If you want to use TFDS-based data, check out other transfer configs. Remember to set `TFDS_DATA_DIR` to point to the folder with data (can be GCP data bucket).
+
+
+## Model Development Contributions
+
+See the [technical report](https://arxiv.org/abs/2407.07726)'s Appendix.

big_vision_repo/big_vision/configs/proj/paligemma/transfers/activitynet_cap.py

@@ -0,0 +1,209 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to ActivityNet Video captioning.
+
+IMPORTANT: This config is based on an unreleased version of DeepMind Video
+Readers (DMVR). Users can either set up DMVR using the open source code from
+GitHub (see below for details), or add their own data loader of choice.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+TEXT_LEN = 64
+DATASET_NAME = 'activitynet_captions_mr'
+# Numbers might need to be updated due to wipeout. Current from 2024-04-28
+SPLIT_SIZE = {'train': 30545, 'valid': 14338, 'test': 13982}
+
+
+def training_data(res, *, final_split, num_frames=8, stride=None):
+  """Creates training data config.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Train on all train+valid data.
+    num_frames: number of sampled frames per video.
+    stride: stride at which the frames are sampled.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  pp = '|'.join([
+      # prepare the frames by decoding, resizing, replicating, sampling:
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      # pick one caption at random during training (there is actually just one!)
+      'strfmt("caption en", outkey="prefix")',
+      'video_choice(inkey="caption/string", outkey="suffix")',
+      combine_and_keep_train(TEXT_LEN),
+  ])
+
+  c = bvcc.parse_arg('')
+  c.data = {}
+  splits = ['train', 'valid'] if final_split else ['train']
+  raise NotImplementedError('Please implement a video reader of choice!')
+  # For example DMVR https://github.com/google-deepmind/dmvr
+  # The reader should support the following arguments:
+  # - name: Name of the reader.
+  # - dataset_name: Name of the data set.
+  # - split: Data set split.
+  # - num_frames: Number of frames sampled from the video.
+  # - stride: Stride at which the video frames are sampled.
+  # - deterministic_fs: Whether to sample the frames starting at the first
+  #   frame or whether an offest should be chosen at random (if there are more
+  #   frames than num_frames * stride)
+  # - first_k_shards: Whether to only use the first k shards of the data
+  #   (optional but useful for speeding up intermediate evaluations).
+  for split in splits:
+    c.data[split] = SPLIT_SIZE[split]
+    c[split] = {'pp': pp}
+    c[split].data = dict(
+        # PLEASE ADD YOUR READER HERE:
+        name='<add_your_data_loader_here>',
+        dataset_name=DATASET_NAME, split=split,
+        num_frames=num_frames, stride=stride,
+        deterministic_fs=False)
+  return c
+
+
+def add_eval(c, res, num_frames=8, stride=None):  # pylint: disable=unused-argument
+  """Captioning evaluator."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  pp = '|'.join([
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      'strfmt("caption en", outkey="prefix")',
+      'strfmt("{example/video_id}[{segment_start}-{segment_end}]", outkey="image/id")',
+      'copy("caption/string", "captions")',
+      combine_and_keep_eval(TEXT_LEN, keep=('image/id', 'captions')),
+  ])
+
+  for freq, name, split, first_k_shards, skip_first_eval in [
+      (1/8, 'minitrain', 'train', 2, False),  # To gauge memorization.
+      (1/4, 'minival', 'valid', 2, False),  # To monitor val progress.
+      (1, 'val', 'valid', None, False),  # To tune hparams.
+      (1, 'eval', 'test', None, False),  # final metric
+  ]:
+    c.evals[f'{DATASET_NAME}/{name}'] = dict(
+        type='proj.paligemma.transfers.coco_caption',
+        pred='decode', pred_kw={'max_decode_len': TEXT_LEN},
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        log_percent=freq, tokenizer=TOKENIZER,
+        pp_fn=pp, skip_first=skip_first_eval)
+
+
+def add_eval_pplx(c, res, num_frames=8, stride=None):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  for name, split, first_k_shards in [
+      ('minitrain', 'train', 2),  # To gauge memorization.
+  ]:
+    c.evals[f'{DATASET_NAME}/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=1/8,  # Not too cheap, do 10x per run.
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        pp_fn=c_train.train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=1e-5, wd=1e-6, total_epochs=1, **bvcc.arg(freeze_vit=True, res=224, **c))
+
+
+sweep = sweep_best
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', num_frames=16, stride=30, res=224,
+                     freeze_vit=False, freeze_llm=False, final_split=False)
+
+  c.input = training_data(
+      c.res, final_split=c.final_split,
+      num_frames=c.num_frames, stride=c.stride)
+
+  c.total_epochs = 3
+  c.input.batch_size = 128
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 3e-6
+  c.wd = 3e-7
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+
+  # Learning-rate schedule.
+  sched = dict(decay_type='cosine', warmup_percent=0.05)
+  c.schedule = [
+      ('img/.*', None if c.freeze_vit else sched),
+      ('llm/.*', None if c.freeze_llm else sched),
+  ]
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, c.num_frames, c.stride)
+  add_eval_pplx(c, c.res, c.num_frames, c.stride)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  for split in c.input.data.keys():
+    c.input[split].shuffle_buffer_size = 10_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops',
+                  'proj.paligemma.video']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    for split in c.input.data.keys():
+      c.input[split].shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.first_k_shards = 1
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('minitrain', 'minival', 'val', 'eval'):
+    m.append(('epoch', f'{DATASET_NAME}/{split}/cider'))
+  for split in ('minitrain', 'minival'):
+    m.append(('epoch', f'{DATASET_NAME}/{split}/pplx/avg'))
+  return m
+

big_vision_repo/big_vision/configs/proj/paligemma/transfers/activitynet_qa.py

@@ -0,0 +1,213 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to ActivityNet Video QA.
+
+IMPORTANT: This config is based on an unreleased version of DeepMind Video
+Readers (DMVR). Users can either set up DMVR using the open source code from
+GitHub (see below for details), or add their own data loader of choice.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+TEXT_LEN = 64
+DATASET_NAME = 'activitynet_qa'
+# Numbers might need to be updated due to wipeout. Current from 2024-04-28
+SPLIT_SIZE = {'train': 27610, 'valid': 15760, 'test': 6900}
+
+
+def training_data(res, *, final_split, num_frames, stride):
+  """Creates training data config.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Train on all train+valid data.
+    num_frames: number of sampled frames per video.
+    stride: stride at which the frames are sampled.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  pp = '|'.join([
+      # prepare the frames by decoding, resizing, replicating, sampling:
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      # only one question/answer per example.
+      'reshape([], key="question")|reshape([], key="answer")',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'copy("answer", "suffix")',
+      combine_and_keep_train(TEXT_LEN),
+  ])
+
+  c = bvcc.parse_arg('')
+  c.data = {}
+  splits = ['train', 'valid'] if final_split else ['train']
+  raise NotImplementedError('Please implement a video reader of choice!')
+  # For example DMVR https://github.com/google-deepmind/dmvr
+  # The reader should support the following arguments:
+  # - name: Name of the reader.
+  # - dataset_name: Name of the data set.
+  # - split: Data set split.
+  # - num_frames: Number of frames sampled from the video.
+  # - stride: Stride at which the video frames are sampled.
+  # - deterministic_fs: Whether to sample the frames starting at the first
+  #   frame or whether an offest should be chosen at random (if there are more
+  #   frames than num_frames * stride)
+  # - first_k_shards: Whether to only use the first k shards of the data
+  #   (optional but useful for speeding up intermediate evaluations).
+  for split in splits:
+    c.data[split] = SPLIT_SIZE[split]
+    c[split] = {'pp': pp}
+    c[split].data = dict(
+        # PLEASE ADD YOUR READER HERE:
+        name='<add_your_data_loader_here>',
+        dataset_name=DATASET_NAME, split=split,
+        num_frames=num_frames, stride=stride,
+        deterministic_fs=False)
+  return c
+
+
+def add_eval(c, res, num_frames, stride):  # pylint: disable=unused-argument
+  """QA evaluator."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  pp = '|'.join([
+      # prepare the frames by decoding, resizing, replicating, sampling:
+      f'video_decode({res})|video_replicate_img({num_frames},{num_frames})',
+      f'video_ensure_shape("image", {(num_frames, res, res, 3)})',
+      # only one question/answer per example.
+      'reshape([], key="question")|reshape([], key="answer")',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'strfmt("{id}#{example/video_id}: {question}", "question_id")',
+      combine_and_keep_eval(TEXT_LEN, keep=('question_id', 'answer')),
+  ])
+
+  for freq, name, split, first_k_shards, skip_first_eval in [
+      (1/8, 'minitrain', 'train', 2, False),  # To gauge memorization.
+      (1/4, 'minival', 'valid', 2, False),  # To monitor val progress.
+      (1, 'val', 'valid', None, True),  # To tune hparams.
+      (1, 'eval', 'test', None, True),  # final metric
+  ]:
+    c.evals[f'activitynet_qa/{name}'] = dict(
+        type='proj.paligemma.transfers.vqa',
+        pred='decode', pred_kw={'max_decode_len': TEXT_LEN},
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        log_percent=freq, tokenizer=TOKENIZER,
+        pp_fn=pp, skip_first=skip_first_eval)
+
+
+def add_eval_pplx(c, res, num_frames, stride):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, num_frames=num_frames, stride=stride)
+
+  for name, split, first_k_shards in [
+      ('minitrain', 'train', 2),  # To gauge memorization.
+      ('minival', 'valid', 2),
+  ]:
+    c.evals[f'activitynet_qa/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=1/8,  # Not too cheap, do 10x per run.
+        data={**c_train.train.data, 'split': split,
+              'first_k_shards': first_k_shards,
+              'deterministic_fs': True},
+        pp_fn=c_train.train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=1e-5, wd=1e-6, total_epochs=1, **bvcc.arg(num_frames=16, stride=70, res=224, **c))
+
+
+sweep = sweep_best
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', num_frames=16, stride=70, res=224,
+                     freeze_vit=False, freeze_llm=False, final_split=False)
+
+  c.input = training_data(
+      c.res, final_split=c.final_split,
+      num_frames=c.num_frames, stride=c.stride)
+
+  c.total_epochs = 3
+  c.input.batch_size = 128
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 1e-5
+  c.wd = 1e-6
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+
+  # Learning-rate schedule.
+  sched = dict(decay_type='cosine', warmup_percent=0.05)
+  c.schedule = [
+      ('img/.*', None if c.freeze_vit else sched),
+      ('llm/.*', None if c.freeze_llm else sched),
+  ]
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, c.num_frames, c.stride)
+  add_eval_pplx(c, c.res, c.num_frames, c.stride)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  for split in c.input.data.keys():
+    c.input[split].shuffle_buffer_size = 10_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops',
+                  'proj.paligemma.video']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    for split in c.input.data.keys():
+      c.input[split].shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.first_k_shards = 1
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('minitrain', 'minival', 'val', 'eval'):
+    m.append(('epoch', f'{DATASET_NAME}/{split}/acc'))
+  for split in ('minitrain', 'minival'):
+    m.append(('epoch', f'{DATASET_NAME}/{split}/pplx/avg'))
+  return m
+

big_vision_repo/big_vision/configs/proj/paligemma/transfers/ai2d.py

@@ -0,0 +1,170 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to AI2D.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+PREFIX = 'answer en '
+PROMPT = 'choose from:'
+PROMPT_SEP = ' \t '
+
+
+def training_data(res, final_split, text_len=128):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: whether to use all train data.
+    text_len: sequence length
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='ai2d',
+      # 12k training examples.
+      split='train' if final_split else 'train[:-1024]',
+  )
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      f'strjoin("{PROMPT_SEP}", inkey="possible_answers", outkey="ansstr")',
+      f'strfmt("{PREFIX} {{question}} {PROMPT} {{ansstr}}", outkey="prefix")',
+      'copy(inkey="answer", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=128, **kw):
+  """AI2D evaluators."""
+  pp = '|'.join([
+      f'decode|resize({res})|value_range(-1, 1)',
+      f'strjoin("{PROMPT_SEP}", inkey="possible_answers", outkey="ansstr")',
+      f'strfmt("{PREFIX} {{question}} {PROMPT} {{ansstr}}", outkey="prefix")',
+      'copy(inkey="id",outkey="question_id")',
+      combine_and_keep_eval(text_len, keep=('answer', 'question_id')),
+  ])
+
+  for name, split in [
+      ('minitrain', 'train[:1024]'),  # To gauge memorization.
+      ('minival', 'train[-1024:]'),   # To tune hparams.
+      ('eval', 'test'),               # To compute final publishable scores.
+  ]:
+    c.evals[f'ai2d/{name}'] = dict(
+        type='proj.paligemma.transfers.vqa',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        outfile=f'{{workdir}}/ai2d_{name}.json',
+        to_lower=False,  # Model sees options in prompt and can match the case.
+        data={**training_data(res, True, text_len).data, 'split': split},
+        log_percent=1/8, tokenizer=TOKENIZER, pp_fn=pp)
+    c.evals[f'ai2d/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=128):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, True, text_len)  # Use mostly same settings as training.
+
+  for name, split in [
+      ('minitrain', 'train[:1024]'),  # To gauge memorization.
+      ('minival', 'train[-1024:]'),   # To tune hparams.
+      ('eval', 'test'),               # To compute final publishable scores.
+  ]:
+    c.evals[f'ai2d/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=1/8,
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=1e-5, wd=1e-6, total_epochs=10, **bvcc.arg(res=224, **c))
+  add(lr=1e-5, wd=1e-6, total_epochs=10, **bvcc.arg(res=448, **c))
+  # 896 was not better than 448 ((internal link)).
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224, final_split=False)
+
+  c.input = training_data(c.res, final_split=c.final_split)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 10
+  c.input.batch_size = 256
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 1e-5
+  c.wd = 1e-5 * 0.1
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=1024)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('eval', 'minival', 'minitrain'):
+    m.append(f'ai2d/{split}/pplx/avg')
+    m.append(f'ai2d/{split}/acc')
+  return m

big_vision_repo/big_vision/configs/proj/paligemma/transfers/aokvqa_da.py

@@ -0,0 +1,161 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to A-OK-VQA using Direct Answer mode.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+def training_data(res, final_split, text_len=32):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Whether to use train and validation data.
+    text_len: sequence length
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='aokvqa',
+      split='train + val' if final_split else 'train',
+  )
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'choice_no_replacement(inkey="direct_answers", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=32, **kw):
+  """We can use the normal VQA evaluators."""
+  pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'copy(inkey="direct_answers", outkey="answers")',
+      combine_and_keep_eval(text_len, keep=('answers', 'question_id')),
+  ])
+
+  for freq, name, split in [
+      (1/4, 'minitrain', 'train[:5%]'),  # To gauge memorization.
+      (1/4, 'eval', 'val'),              # To tune hparams.
+      (1.0, 'test', 'test'),             # To compute final predictions.
+  ]:
+    c.evals[f'aokvqa_da/{name}'] = dict(
+        type='proj.paligemma.transfers.vqa',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        outfile=f'{{workdir}}/aokvqa_da_{name}.json',
+        data={**training_data(res, True, text_len).data, 'split': split},
+        log_percent=freq, skip_first=freq == 1, tokenizer=TOKENIZER, pp_fn=pp)
+    c.evals[f'aokvqa/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=32):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, True, text_len)  # Use mostly same settings as training.
+
+  for name, split in [
+      ('minitrain', 'train[:5%]'),  # To gauge memorization.
+      ('eval', 'val'),              # To tune hparams.
+  ]:
+    c.evals[f'aokvqa_da/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=0.05,  # Eval ~20x per run; it's cheap.
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=5e-6, wd=0.0, **bvcc.arg(res=224, **c))
+  add(lr=5e-6, wd=0.0, **bvcc.arg(res=448, **c))
+  # not better: add(lr=5e-6, wd=0.0, **bvcc.arg(res=896, **c))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224, final_split=False)
+
+  c.input = training_data(c.res, c.final_split)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 10
+  c.input.batch_size = 128
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 5e-6
+  c.wd = 0.0
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=256)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('eval', 'minival', 'minitrain'):
+    m.append(f'aokvqa/{split}/pplx/avg')
+    m.append(f'aokvqa/{split}/acc')
+  return m

big_vision_repo/big_vision/configs/proj/paligemma/transfers/aokvqa_mc.py

@@ -0,0 +1,169 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to A-OK-VQA using multiple choice answers.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+PREFIX = 'answer en '
+PROMPT = 'choose from:'
+PROMPT_SEP = ' \t '
+
+
+def training_data(res, final_split, text_len=128):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Whether to use train and validation data.
+    text_len: sequence length
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='aokvqa',
+      split='train + val' if final_split else 'train',
+  )
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      f'strjoin("{PROMPT_SEP}", inkey="multiple_choice_possible_answers", outkey="ansstr")',
+      f'strfmt("{PREFIX} {{question}} {PROMPT} {{ansstr}}", outkey="prefix")',
+      'getidx(inkey="multiple_choice_possible_answers", index_key="multiple_choice_correct_idx", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=128, **kw):
+  """VQAv2 evaluators."""
+  pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      f'strjoin("{PROMPT_SEP}", inkey="multiple_choice_possible_answers", outkey="ansstr")',
+      f'strfmt("{PREFIX} {{question}} {PROMPT} {{ansstr}}", outkey="prefix")',
+      'getidx(inkey="multiple_choice_possible_answers", index_key="multiple_choice_correct_idx", outkey="answer")',
+      combine_and_keep_eval(text_len, keep=('answer', 'question_id')),
+  ])
+
+  for freq, name, split in [
+      (1/4, 'minitrain', 'train[:5%]'),  # To gauge memorization.
+      (1/4, 'eval', 'val'),              # To tune hparams.
+      (1.0, 'test', 'test'),             # To compute final predictions.
+  ]:
+    c.evals[f'aokvqa_mc/{name}'] = dict(
+        type='proj.paligemma.transfers.vqa',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        outfile=f'{{workdir}}/aokvqa_mc_{name}.json',
+        data={**training_data(res, True, text_len).data, 'split': split},
+        log_percent=freq, skip_first=freq == 1, tokenizer=TOKENIZER, pp_fn=pp)
+    c.evals[f'aokvqa/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=128):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, True, text_len)  # Use mostly same settings as training.
+
+  for name, split in [
+      ('minitrain', 'train[:5%]'),  # To gauge memorization.
+      ('eval', 'val'),              # To tune hparams.
+      ('test', 'test'),             # To compute final predictions.
+  ]:
+    c.evals[f'aokvqa_mc/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=0.05,  # Eval ~20x per run; it's cheap.
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  add(lr=5e-6, wd=0.0, **bvcc.arg(res=224, **c))
+  add(lr=5e-6, wd=0.0, **bvcc.arg(res=448, **c))
+  # add(lr=5e-6, wd=0.0, **bvcc.arg(res=896, **c))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224, final_split=False)
+
+  c.input = training_data(c.res, c.final_split)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 15
+  c.input.batch_size = 128
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 5e-6
+  c.wd = 0.0
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=256)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('eval', 'minival', 'minitrain'):
+    m.append(f'aokvqa/{split}/pplx/avg')
+    m.append(f'aokvqa/{split}/acc')
+  return m