jit transition (#53)

Author: akolesnikoff

README.md

@@ -20,6 +20,12 @@ 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:
 
@@ -51,6 +57,13 @@ codebase:
   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.
 
 ### Multimodal research
 
@@ -64,21 +77,28 @@ codebase:
 - [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*.
+- [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.
 
-### Knowledge distillation
+### 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).
-
-### Training
-
 - [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
 
@@ -118,7 +138,7 @@ details, but generally speaking, running on a GPU machine involves calling
 `python -m COMMAND` while running on TPUs, including multi-host, involves
 
 ```
-gcloud alpha compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all
+gcloud compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all
   --command "bash big_vision/run_tpu.sh COMMAND"
 ```
 
@@ -273,7 +293,7 @@ The following command line will create TPU VMs with 32 cores,
 4 hosts.
 
 ```
-gcloud alpha compute tpus tpu-vm create $NAME --zone $ZONE --accelerator-type v3-32 --version v2-tf-stable
+gcloud compute tpus tpu-vm create $NAME --zone $ZONE --accelerator-type v3-32 --version tpu-ubuntu2204-base
 ```
 
 ## Install `big_vision` on TPU VMs
@@ -283,8 +303,8 @@ dependencies.
 
 ```
 git clone https://github.com/google-research/big_vision
-gcloud alpha compute tpus tpu-vm scp --recurse big_vision/big_vision $NAME: --zone=$ZONE --worker=all
-gcloud alpha compute tpus tpu-vm ssh $NAME --zone=$ZONE --worker=all --command "bash big_vision/run_tpu.sh"
+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
@@ -298,13 +318,13 @@ Specifically, the seven TFDS datasets used during evaluations will be generated
 under `~/tensorflow_datasets` on TPU machine with this command:
 
 ```
-gcloud alpha 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"
+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 alpha 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"
+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,
@@ -322,23 +342,28 @@ The following command line fine-tunes a pre-trained `vit-i21k-augreg-b/32` model
 on `cifar10` dataset.
 
 ```
-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.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"
+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"
 ```
 
+## Checkpointing on cloud
+
+In the past, we recommended writing checkpoints to a Google Cloud Bucket. With the latest update, this is very slow because of technical issues with the checkpointing format.
+We are working on a solution, but in the meantime, we have updated our instructions to write checkpoints to a local folder on the TPU machine. Don't forget to copy useful checkpoints elsewhere after training.
+
 ## 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 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.train --config big_vision/configs/bit_i1k.py  --workdir gs://$GS_BUCKET_NAME/big_vision/workdir/`date '+%m-%d_%H%M'`"
+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'`"
 ```
 
 ## Sometimes useful gcloud commands
 
-- Destroy the TPU machines: `gcloud alpha compute tpus tpu-vm delete $NAME --zone $ZONE`
-- Remove all big_vision-related folders on all hosts: `gcloud alpha compute tpus tpu-vm ssh $NAME --zone $ZONE --worker=all --command 'rm -rf ~/big_vision ~/bv_venv'`
+- 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'`
 
 # ViT baseline
 

big_vision/configs/bit_i1k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.

big_vision/configs/bit_i21k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.

big_vision/configs/common.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -130,9 +130,16 @@ def autotype(x):
 
 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]:

big_vision/configs/common_fewshot.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -18,12 +18,20 @@
 
 
 def get_fewshot_lsr(target_resolution=224, resize_resolution=256,
-                    runlocal=False):
+                    runlocal=False, **kw):
   """Returns a standard-ish fewshot eval configuration."""
-  config = mlc.ConfigDict()
+  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.representation_layer = 'pre_logits'
-  config.log_steps = 25_000
   config.datasets = {
       'caltech': ('caltech101', 'train', 'test'),  # copybara:srtip
       'cars': ('cars196:2.1.0', 'train', 'test'),
@@ -37,12 +45,12 @@ def get_fewshot_lsr(target_resolution=224, resize_resolution=256,
   } if not runlocal else {
       'pets': ('oxford_iiit_pet', 'train', 'test'),
   }
-  config.pp_train = f'decode|resize({resize_resolution})|central_crop({target_resolution})|value_range(-1,1)|keep("image", "label")'
-  config.pp_eval = f'decode|resize({resize_resolution})|central_crop({target_resolution})|value_range(-1,1)|keep("image", "label")'
-  config.shots = (1, 5, 10, 25)
-  config.l2_reg = 2.0 ** 10
-  config.num_seeds = 3
+  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)]
-  config.prefix = ''  # No prefix as we do already prefix with a/ z/ and zz/
 
   return config

big_vision/configs/load_and_eval.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -36,25 +36,35 @@
 
 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 lit_eval
-import ml_collections as mlc
+# from big_vision.configs.proj.image_text import lit_eval
+
+
+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
+  return config
 
 
 def get_config(arg='name=bit_paper,batch_size=2'):
-  arg = bvcc.parse_arg(arg, name='', batch_size=2)
-  config = mlc.ConfigDict()
-  config.batch_size_eval = arg.batch_size
+  config = bvcc.parse_arg(arg, name='', batch_size=2)
+
+  # 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'),
+  ))
 
   # 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()[arg.name](config)
+  globals()[config.name](config)
   return config
 
 
 def bit_paper(config):
-  # We could omit init_{shapes,types} if we wanted, as they are the default.
-  config.init_shapes = [(1, 224, 224, 3)]
-  config.init_types = ['float32']
   config.num_classes = 1000
 
   config.model_name = 'bit_paper'
@@ -82,9 +92,6 @@ def get_eval(split, lbl, dataset='imagenet2012_real'):
 
 
 def vit_i1k(config):
-  # We could omit init_{shapes,types} if we wanted, as they are the default.
-  config.init_shapes = [(1, 224, 224, 3)]
-  config.init_types = ['float32']
   config.num_classes = 1000
 
   config.model_name = 'vit'
@@ -104,9 +111,6 @@ def vit_i1k(config):
 
 
 def mlp_mixer_i1k(config):
-  # We could omit init_{shapes,types} if we wanted, as they are the default.
-  config.init_shapes = [(1, 224, 224, 3)]
-  config.init_types = ['float32']
   config.num_classes = 1000
 
   config.model_name = 'mlp_mixer'
@@ -125,9 +129,6 @@ def mlp_mixer_i1k(config):
 
 
 def vit_i21k(config):
-  # We could omit init_{shapes,types} if we wanted, as they are the default.
-  config.init_shapes = [(1, 224, 224, 3)]
-  config.init_types = ['float32']
   config.num_classes = 21843
 
   config.model_name = 'vit'

big_vision/configs/mlp_mixer_i1k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.

big_vision/configs/transfer.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -148,11 +148,11 @@ def _set_imagenet_variants(config, h_res=448, l_res=384):
 def get_config(arg=None):
   """Config for adaptation."""
   arg = bvcc.parse_arg(arg, model='vit', dataset='cifar10', crop='resmall_crop',
-                       h_res=448, l_res=384, runlocal=False)
+                       h_res=448, l_res=384, batch_size=512, runlocal=False)
   config = mlc.ConfigDict()
 
   config.input = {}
-  config.input.batch_size = 512 if not arg.runlocal else 8
+  config.input.batch_size = arg.batch_size if not arg.runlocal else 8
   config.input.shuffle_buffer_size = 50_000 if not arg.runlocal else 100
 
   config.log_training_steps = 10

big_vision/configs/vit_i1k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -15,6 +15,9 @@
 # pylint: disable=line-too-long
 r"""Pre-training ViT on ILSVRC-2012 as in https://arxiv.org/abs/2106.10270
 
+This config does NOT include regularization (dropout, stochastic depth), which
+was shown to help with B/32, B/16, L/16 models in the paper (Figure 4).
+
 This configuration makes use of the "arg" to get_config to select which model
 to run, so a few examples are given below:
 

big_vision/configs/vit_i21k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.

big_vision/configs/vit_s16_i1k.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.

big_vision/datasets/core.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.
@@ -21,7 +21,7 @@
 class DataSource:
   """The API that any data source should implement."""
 
-  def get_tfdata(self, ordered):
+  def get_tfdata(self, ordered, *, process_split=True):
     """Creates this data object as a tf.data.Dataset.
 
     This will be called separately in each process, and it is up to the dataset
@@ -30,6 +30,8 @@ def get_tfdata(self, ordered):
     Args:
       ordered: if True, the dataset should use deterministic ordering, if False
         it may have undefined ordering. Think of True == val, False == train.
+      process_split: if False then every process receives the entire dataset
+        (e.g.  for evaluators running in a single process).
 
     Returns:
       A tf.data.Dataset object.

big_vision/datasets/imagenet/class_names.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Big Vision Authors.
+# 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.