added pali inference

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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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/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

big_vision/configs/proj/paligemma/transfers/chartqa.py

@@ -0,0 +1,180 @@
+# 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 chartqa.
+"""
+
+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
+
+_DATASETS = ('chartqa/human', 'chartqa/augmented')
+# We use the true dataset sizes from https://arxiv.org/pdf/2203.10244.pdf.
+_WEIGHTS = (7_398, 20_901)
+
+
+def training_data(res, *, final_split=False, text_len=48):
+  """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: Train on all train+val data.
+    text_len: sequence length.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'copy(inkey="question", outkey="prefix")',
+      'copy(inkey="answer", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  c.data = {ds: weight for ds, weight in zip(_DATASETS, _WEIGHTS)}
+  for ds in c.data:
+    c[ds] = dict(
+        shuffle_buffer_size=50_000,
+        pp=pp,
+        data=dict(
+            name=ds,
+            split='train+val' if final_split else 'train',
+        ),
+    )
+  return c
+
+
+def add_eval(c, res, text_len=48, **kw):
+  """Add eval configs."""
+  c_train = training_data(res, final_split=True, text_len=text_len)
+
+  pp_eval = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'copy(inkey="question", outkey="prefix")',
+      combine_and_keep_eval(text_len, keep=('answer', 'question_id')),
+  ])
+
+  for name, split in [
+      ('minitrain', 'train[:5%]'),
+      ('minival', 'val'),
+      ('eval', 'test'),
+  ]:
+    for ds in _DATASETS:
+      c.evals[f'{ds}/{name}'] = dict(
+          type='proj.paligemma.transfers.chartqa',
+          pred='decode', pred_kw={'max_decode_len': text_len},
+          to_lower=True,
+          outfile=f'{{workdir}}/{ds.replace("/", "_")}_{name}.json',
+          data={**c_train[ds].data, 'split': split},
+          log_percent=0.1, tokenizer=TOKENIZER, pp_fn=pp_eval)
+      c.evals[f'{ds}/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=48):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, text_len=text_len)  # Use mostly same settings as training.
+  for name, split in [
+      ('minitrain', 'train[:5%]'),  # To gauge memorization.
+      ('minival', 'val'),  # To tune hparams.
+      ('eval', 'test'),  # To compute final publishable scores.
+  ]:
+    for ds in _DATASETS:
+      c.evals[f'{ds}/{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[ds].data, 'split': split},
+          pp_fn=c_train[ds].pp,
+      )
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  # TODO: Update once latest numbers are in and have only 1 setup.
+  # Based on (internal link) (relaxed_accuracy).
+  add(lr=1e-5, wd=1e-6, total_epochs=30, **bvcc.arg(res=224, **c))
+  # Based on sweep (internal link) and on (internal link) (relaxed_accuracy).
+  add(lr=1e-5, wd=1e-6, total_epochs=30, **bvcc.arg(res=448, **c))
+  # Based on (internal link) (relaxed_accuracy).
+  # Not better: add(lr=1e-5, wd=1e-6, total_epochs=30, **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=896, 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 = 30
+  c.input.batch_size = 256
+  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.2
+  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.1)
+  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.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'):
+    for ds in _DATASETS:
+      c.input[ds].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'):
+    for ds in _DATASETS:
+      m.append(f'{ds}/{split}/relaxed_acc')
+      m.append(f'{ds}/{split}/pplx/avg')
+  return m