Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Tipsomaly
+emoji: ⚡
+colorFrom: gray
+colorTo: green
+sdk: gradio
+sdk_version: 6.6.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: TIPS Over Tricks for Zero-shot Anomaly Detection
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

Tipsomaly/.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.pyc
+*.jpg
+*.png

Tipsomaly/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Alireza Salehi et al.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights  
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell      
+copies of the Software, and to permit persons to whom the Software is         
+furnished to do so, subject to the following conditions:                       
+
+The above copyright notice and this permission notice shall be included in    
+all copies or substantial portions of the Software.                           
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR    
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,      
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE   
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER        
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN     
+THE SOFTWARE.

Tipsomaly/README.md

@@ -0,0 +1,131 @@
+# Tipsomaly (ICASSP 2026)
+
+Official PyTorch implementation of [TIPS Over Tricks: Simple Prompts for Effective Zero-shot Anomaly Detection](https://arxiv.org/abs/2602.03594) — a spatially-aware zero-shot anomaly detection pipeline built on the [TIPS](https://arxiv.org/abs/2410.16512) vision-language model, using decoupled prompts and local evidence injection to improve image-level and pixel-level performance.
+
+---
+
+## Table of Contents
+
+- [📖 Introduction](#-introduction)
+- [📊 Results](#-results)
+- [🚀 Quickstart](#-quickstart)
+- [📍 Checkpoints](#-checkpoints)
+- [🔧 Setup](#-setup)
+- [🗂️ Datasets](#️-datasets)
+- [🛠️ Training](#️-training)
+- [⚖️ License](#️-license)
+- [🔗 Citation](#-citation)
+<!-- - [🙏 Acknowledgements](#-acknowledgements) -->
+
+---
+
+## 📖 Introduction
+Anomaly detection identifies departures from expected behavior in safety-critical settings. When target-domain normal data are unavailable, zero-shot anomaly detection (ZSAD) leverages vision-language models (VLMs). However, CLIP's coarse image-text alignment limits both localization and detection due to (i) spatial misalignment and (ii) weak sensitivity to fine-grained anomalies; prior work compensates with complex auxiliary modules yet largely overlooks the choice of backbone. We revisit the backbone and use TIPS-a VLM trained with spatially aware objectives. While TIPS alleviates CLIP's issues, it exposes a distributional gap between global and local features. We address this with decoupled prompts-fixed for image-level detection and learnable for pixel-level localization-and by injecting local evidence into the global score. Without CLIP-specific tricks, our TIPS-based pipeline improves image-level performance by 1.1-3.9% and pixel-level by 1.5-6.9% across seven industrial datasets, delivering strong generalization with a lean architecture.
+
+![Architecture](imgs/Models_Architecture_page-0001.jpg)
+
+---
+
+## 📊 Results
+Across 14 industrial and medical benchmarks, Tipsomaly consistently outperforms prior CLIP-based zero-shot methods while remaining lightweight. The figure below summarizes its performance across datasets.
+
+![Results](imgs/results-table.png)
+
+We compare pixel-level anomaly maps from Tipsomaly with prior CLIP-based methods (AdaCLIP and AnomalyCLIP) on industrial and medical samples. As shown in the following figure, Tipsomaly more accurately localizes anomalous regions across both domains.
+
+![Qualitative-results](imgs/Qualitative_results_page-0001.jpg)
+
+---
+
+## 🚀 Quickstart
+**Quick start (recommended):**  
+Use our ready-to-run [Kaggle notebook](https://www.kaggle.com/code/sepehrnoey/tipsomaly-reproduction) to reproduce the results with the provided checkpoints.
+
+## 📍 Checkpoints
+You can find our checkpoints trained on MVTec and VisA in the directory [workspaces](./workspaces/) or on google drive at [this link](https://drive.google.com/file/d/1yvgZYHFljFkGwTD0DYOdrZ_95ZbmjDcC/view?usp=sharing).
+
+---
+
+
+
+
+## 🔧 Setup
+To setup the environment and run experiments, follow these steps.
+1. Install the required dependencies with the following command.
+```
+pip install -r requirements.txt
+```
+2. Download your desired checkpoints of the TIPS model. For this purpose, edit and run the mentioned script to download model components. After downloading the components, place them in a desired directory.
+```
+bash model/tips/checkpoints/download_checkpoints.sh
+```
+```
+mkdir tips && mv tips_oss_l14_highres_distilled_text.npz tips_oss_l14_highres_distilled_vision.npz tokenizer.model tips/
+```
+3. Prepare your datasets in the format described in the [Datasets](#️-datasets) section and place them in a desired directory.
+```
+mkdir data-root && mv mvtec visa data-root/
+``` 
+4. Now, by setting the correct paths for the model, checkpoints, and the data root path, you can edit the runnable script `reproduce.sh` to run experiments. Here is a sample:
+```
+models_dir="path/to/tips"
+data_root_dir="path/to/data-root"
+model_version='l14h' 
+checkpoint_path="./workspaces/trained_on_mvtec_default/vegan-arkansas/checkpoints"
+
+python test.py --models_dir $models_dir --checkpoint_path $checkpoint_path --data_root_dir $data_root_dir --dataset visa --epoch 2 --model_version $model_version --fixed_prompt_type industrial
+
+```
+
+## 🗂️ Datasets
+The datasets should be set in the following format. Each dataset, must contain a file named `meta.json` which contains information about the dataset. Common datasets for anomaly detection and instructions to generate `meta.json` files can be found in the [AdaCLIP GitHub repository](https://github.com/caoyunkang/AdaCLIP). The general expected structure of the datasets is like the following:
+```
+data-root/
+├── mvtec/
+│   ├── meta.json
+│   ├── bottle/
+│   │   ├── train/
+│   │   │   └── good/
+│   │   │       └── 000.png
+│   │   ├── test/
+│   │   │   ├── good/
+│   │   │   │   └── 000.png
+│   │   │   ���── broken_large/
+│   │   │   │   └── 000.png
+│   │   │   └── broken_small/
+│   │   │       └── 000.png
+│   │   └── ground_truth/
+│   │       ├── broken_large/
+│   │       │   └── 000_mask.png
+│   │       └── broken_small/
+│   │           └── 000_mask.png
+│
+├── visa/
+├── mpdd/
+├── ....
+....
+```
+
+## 🛠️ Training  
+You can use our ready-to-run [Kaggle notebook](https://www.kaggle.com/code/sepehrnoey/tipsomaly-train) to train your model and test on your desired dataset. Steps to prepare the datasets and downloading the base model components can be found at the [Setup](#-setup) section. You can also find the commands to train the model and test multiple datasets in a loop in the [train_test.sh](train_test.sh)
+
+## 🔗 Citation
+If you find this project helpful for your research, please consider citing the following BibTeX entry.
+
+
+<!-- 📚 [Paper Link](https://arxiv.org/pdf/2504.11055) -->
+
+**BibTeX:**
+```bibtex
+@article{salehi2026tips,
+  title={TIPS Over Tricks: Simple Prompts for Effective Zero-shot Anomaly Detection},
+  author={Salehi, Alireza and Karami, Ehsan and Noey, Sepehr and Noey, Sahand and Yamada, Makoto and Hosseini, Reshad and Sabokrou, Mohammad},
+  journal={arXiv preprint arXiv:2602.03594},
+  year={2026}
+}
+```
+
+<!-- ## 🙏 Acknowledgements -->
+
+## ⚖️ License
+This project is licensed under the MIT License. Please refer to the [LICENSE](LICENSE) file for more information.

Tipsomaly/datasets/dataset.py

@@ -0,0 +1,115 @@
+import torch.utils.data as data
+import json
+import random
+from PIL import Image
+import numpy as np
+import os
+
+def sample_data(meta_info, k_shot):
+    sampled_data = []
+    complement_data = [] 
+    
+    for cls_name, data_list in meta_info.items():
+        nrm_smpls = [item for item in data_list if item['anomaly'] == 0]
+        anm_smpls = [item for item in data_list if item['anomaly'] == 1]
+        
+        if k_shot > 0:
+            n_samples = k_shot
+            n_nrm_smpls = min(int(n_samples / 2), len(nrm_smpls))
+            n_anm_smpls = min(int(n_samples / 2), len(anm_smpls))
+
+        cls_data = []
+        cls_data.extend(random.sample(nrm_smpls, n_nrm_smpls))
+        cls_data.extend(random.sample(anm_smpls, n_anm_smpls))
+        sampled_data.extend(cls_data)
+        
+        complement_class_data = [item for item in data_list if item not in cls_data]
+        complement_data.extend(complement_class_data)
+        
+        print(f'num samples for cls {cls_name}, norm: {n_nrm_smpls}, anom: {n_anm_smpls}')
+        
+    return sampled_data, complement_data
+
+class Dataset(data.Dataset):
+    def __init__(self, roots, transform, target_transform, kwargs=None):
+        self.roots = roots
+        self.transform = transform
+        self.target_transform = target_transform
+        split='test'
+        
+        meta_infos = {}
+        for root in roots:
+            with open(f'{root}/meta.json', 'r') as f:
+                meta_info = json.load(f)
+                for cls in meta_info[split]:
+                    meta_info[split][cls] = [{**s, 'root': root} for s in meta_info[split][cls]]
+                    
+                    if cls in meta_infos:
+                        meta_infos[cls].extend(meta_info[split][cls])
+                        meta_infos[cls].extend(meta_info[split][cls])
+                    else:
+                        meta_infos[cls] = meta_info[split][cls]
+                        
+        meta_info_classes = list(meta_infos.keys())
+
+        self.selected_class = kwargs.class_name
+        self.cls_names = meta_info_classes if self.selected_class == ['all'] else self.selected_class
+        
+        self.data_all = []
+        for cls_name in self.cls_names:
+            self.data_all.extend(meta_infos[cls_name])
+
+        self.dataset_name = kwargs.dataset
+        self.class_ids = list(range(len(self.cls_names)))
+        self.class_name_map_class_id = {k: index for k, index in zip(self.cls_names, self.class_ids)}
+
+        # Few-shot dataset (splitting...)
+        self.k_shot = kwargs.k_shot
+        if not self.k_shot == 0:
+            sampled_sets = sample_data(meta_infos, self.k_shot)
+            
+            if kwargs.type == 'train':
+                self.data_all = sampled_sets[0] 
+            
+            elif kwargs.type == 'test' and kwargs.train_dataset == self.dataset_name:
+                self.data_all = sampled_sets[1] 
+        
+        self.length = len(self.data_all)
+        print(f"number of dataset samples: {self.length}")
+            
+    def _process_image(self, data):
+        img_path, mask_path, cls_name, specie_name, anomaly = data['img_path'], data['mask_path'], \
+                                                                data['cls_name'], data['specie_name'], data['anomaly']
+        
+        root = data['root']    
+        full_img_path = os.path.join(root, img_path)
+        img = Image.open(full_img_path).convert('RGB')  
+
+        if anomaly == 0 or (not os.path.isfile(os.path.join(root, mask_path))):
+            img_mask = Image.fromarray(np.zeros((img.size[1], img.size[0])), mode='L')
+        else:
+            img_mask = np.array(Image.open(os.path.join(root, mask_path)).convert('L')) > 0
+            img_mask = Image.fromarray(img_mask.astype(np.uint8) * 255, mode='L')
+
+        img = self.transform(img) if self.transform is not None else img
+        img_mask = self.target_transform(img_mask)
+        img_mask = np.where(img_mask > 0.5, 1, 0)
+
+        result = {
+            'img': img,
+            'abnorm_mask': img_mask,
+            'cls_name': cls_name,
+            'anomaly': anomaly,
+            'img_path': os.path.join(root, img_path),
+            "cls_id": self.class_name_map_class_id[cls_name]
+        }
+
+        return result
+
+    def __len__(self):
+        return self.length
+    
+    def __getitem__(self, index):
+        data = self.data_all[index]
+        result = self._process_image(data)
+        return result

Tipsomaly/datasets/input_transforms.py

@@ -0,0 +1,23 @@
+from torchvision import transforms
+
+IMAGE_MEAN = (0, 0, 0)
+IMAGE_STD = (1.0, 1.0, 1.0)
+
+class Ensure3Channels:
+    def __call__(self, img):
+        return img.convert('RGB')  
+
+def create_transforms_tips(image_size):
+    transform = transforms.Compose([
+        Ensure3Channels(),
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(IMAGE_MEAN, IMAGE_STD),
+    ])
+
+    target_transform = transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+    ])
+
+    return transform, target_transform

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

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

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

Tipsomaly/model/big_vision/configs/common_fewshot.py

@@ -0,0 +1,60 @@
+# 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, pp=None, **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'),
+  }
+
+  pp = pp or '|'.join([
+      'decode',
+      f'resize({resize_resolution})',
+      f'central_crop({target_resolution})',
+      'value_range(-1,1)'
+  ])
+  pp += '|keep("image", "label")'
+  config.pp_train = pp
+  config.pp_eval = pp
+  config.display_first = [('imagenet', 10)] if not runlocal else [('pets', 10)]
+
+  return config

Tipsomaly/model/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.
+  )

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

Tipsomaly/model/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}
+}
+```

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

Tipsomaly/model/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}
+}
+```

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

Tipsomaly/model/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).

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

Tipsomaly/model/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],
+      )
+  ])

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

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

Tipsomaly/model/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`.

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

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

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

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

Tipsomaly/model/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}
+}
+```

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

Tipsomaly/model/big_vision/configs/proj/givt/givt_demo_colab.ipynb

@@ -0,0 +1,318 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "botgo-GZiWI_"
+      },
+      "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._"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jQxc9UZ-mVrQ"
+      },
+      "source": [
+        "### `big_vision` setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZAXiVta3n2jL"
+      },
+      "outputs": [],
+      "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 \u0026\u0026 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\" \u003e big_vision/big_vision/constraints.txt\n",
+        "!pip install -r big_vision/big_vision/requirements.txt -c big_vision/big_vision/constraints.txt\n",
+        "\n",
+        "# To fix/accelerate gsutil cp\n",
+        "!pip3 -q install --no-cache-dir -U crcmod\n",
+        "\n",
+        "%cd big_vision"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qYS7JNups4MU"
+      },
+      "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",
+      "metadata": {
+        "id": "MaCM_PIcd2Rb"
+      },
+      "source": [
+        "### Select and download model\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7l6QIjdyN3dg"
+      },
+      "outputs": [],
+      "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",
+        "# Manually download VAE checkpoint as the built-in loading logic may hit\n",
+        "# permission issues.\n",
+        "vae_ckpt_local = f\"/tmp/{config.vae.model_init.split('/')[-1]}\"\n",
+        "!test -f {vae_ckpt_local} || gsutil cp {config.vae.model_init} /tmp/\n",
+        "config.vae.model_init = vae_ckpt_local\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])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "SUj5k1bxd6wr"
+      },
+      "source": [
+        "### VAE encode/decode and sampling loop"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vSn7Si2FS1zi"
+      },
+      "outputs": [],
+      "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"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tOnWaJZVeOIX"
+      },
+      "source": [
+        "### Generate images for class label"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "_CiyXD_6nQbu"
+      },
+      "outputs": [],
+      "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]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sCcGB0m1oQY1"
+      },
+      "outputs": [],
+      "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)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "4FWgfAghuh8P"
+      },
+      "outputs": [],
+      "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()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zGPPeXONy0Am"
+      },
+      "outputs": [],
+      "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": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

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

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

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

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

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

Tipsomaly/model/big_vision/configs/proj/image_text/README.md

@@ -0,0 +1,7 @@
+# Image/text models
+
+This directory provides configs and Colabs for different projects on image/text multimodal learning. Please refer to the separate readmes for information on specific projects.
+
+**LiT: Zero-Shot Transfer with Locked-image text Tuning: [README_lit.md](README_lit.md)**
+
+**SigLIP 2: Multilingual Vision-Language Encoders with Improved Semantic Understanding, Localization, and Dense Features: [README_siglip2.md](README_siglip2.md)**

Tipsomaly/model/big_vision/configs/proj/image_text/README_lit.md

@@ -0,0 +1,63 @@
+# 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.

Tipsomaly/model/big_vision/configs/proj/image_text/README_siglip2.md

@@ -0,0 +1,57 @@
+# SigLIP 2: Multilingual Vision-Language Encoders with Improved Semantic Understanding, Localization, and Dense Features
+
+*by Michael Tschannen\*, Alexey Gritsenko\*, Xiao Wang\*, Muhammad Ferjad Naeem\*, Ibrahim Alabdulmohsin\*,Nikhil Parthasarathy\*, Talfan Evans\*, Lucas Beyer\*, Ye Xia, Basil Mustafa, Olivier Hénaff, Jeremiah Harmsen, Andreas Steiner, Xiaohua Zhai\* (\*core contributor)*
+
+[[arxiv]](https://arxiv.org/abs/2502.14786) [[colab]](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP2_demo.ipynb) [[Hugging Face]](https://huggingface.co/collections/google/siglip2-67b5dcef38c175486e240107)
+
+### Summary
+
+We introduce SigLIP 2, a family of new multilingual vision-language encoders that build on the success of the original [SigLIP](https://arxiv.org/abs/2303.15343). In this second iteration, we extend the original image-text training objective with several prior, independently developed techniques into a unified recipe---this includes captioning-based pretraining, self-supervised losses (self-distillation, masked prediction) and online data curation. With these changes, SigLIP 2 models outperform their SigLIP counterparts at all model scales in core capabilities, including zero-shot classification, image-text retrieval, and transfer performance when extracting visual representations for Vision-Language Models (VLMs). Furthermore, the new training recipe leads to significant improvements on localization and dense prediction tasks. We also train variants which support multiple resolutions and preserve the input's native aspect ratio. Finally, we train on a more diverse data-mixture that includes de-biasing techniques, leading to much better multilingual understanding and improved fairness. To allow users to trade off inference cost with performance, we release model checkpoints at four sizes: ViT-B (86M), L (303M), So400m (400M), and g (1B).
+
+### Checkpoints
+
+Below we provide links to all available checkpoints. The standard (non-NaFlex) checkpoints are compatible with [vit.py](https://github.com/google-research/big_vision/blob/main/big_vision/models/vit.py) and [two_towers.py](https://github.com/google-research/big_vision/blob/main/big_vision/models/proj/image_text/two_towers.py) from [SigLIP](https://arxiv.org/abs/2303.15343). The only difference is the vocab size (256k) and the tokenizer (Gemma tokenizer). The NaFlex variant requires a different ViT implementation, [naflex_vit.py](https://github.com/google-research/big_vision/blob/main/big_vision/models/proj/image_text/naflex_vit.py), and an adapted image preprocessing, see [ops_naflex.py](https://github.com/google-research/big_vision/blob/main/big_vision/pp/proj/image_text/ops_naflex.py). The [demo colab](https://colab.research.google.com/github/google-research/big_vision/blob/main/big_vision/configs/proj/image_text/SigLIP2_demo.ipynb) is a good entry point to see how to use the models.
+
+| Model             | ViT       |   Res | Download                                                                              |   INet 0-shot |   COCO  T→I |   COCO  I→T |
+|:------------------|:----------|:------:|:--------------------------------------------------------------------------------------:|:--------------:|:------------:|:------------:|
+| SigLIP 2          | B/32      |   256 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b32_256.npz)         |          74.0 |        47.2 |        63.7 |
+| SigLIP 2          | B/16      |   224 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b16_224.npz)         |          78.2 |        52.1 |        68.9 |
+| SigLIP 2          | B/16      |   256 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b16_256.npz)         |          79.1 |        53.2 |        69.7 |
+| SigLIP 2          | B/16      |   384 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b16_384.npz)         |          80.6 |        54.6 |        71.4 |
+| SigLIP 2          | B/16      |   512 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b16_512.npz)         |          81.2 |        55.2 |        71.2 |
+| SigLIP 2          | L/16      |   256 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_l16_256.npz)         |          82.5 |        54.7 |        71.5 |
+| SigLIP 2          | L/16      |   384 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_l16_384.npz)         |          83.1 |        55.3 |        71.4 |
+| SigLIP 2          | L/16      |   512 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_l16_512.npz)         |          83.5 |        55.2 |        72.1 |
+| SigLIP 2          | So400m/14 |   224 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m14_224.npz)    |          83.2 |        55.1 |        71.5 |
+| SigLIP 2          | So400m/14 |   384 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m14_384.npz)    |          84.1 |        55.8 |        71.7 |
+| SigLIP 2          | So400m/16 |   256 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m16_256.npz)    |          83.4 |        55.4 |        71.5 |
+| SigLIP 2          | So400m/16 |   384 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m16_384.npz)    |          84.1 |        56.0 |        71.2 |
+| SigLIP 2          | So400m/16 |   512 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m16_512.npz)    |          84.3 |        56.0 |        71.3 |
+| SigLIP 2          | g-opt/16  |   256 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_g-opt16_256.npz)     |          84.5 |        55.7 |        72.5 |
+| SigLIP 2          | g-opt/16  |   384 | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_g-opt16_384.npz)     |          85.0 |        56.1 |        72.8 |
+| SigLIP 2 (NaFlex) | B/16      |   var. | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_b16_naflex.npz)      |          78.5 |        51.1 |        67.3 |
+| SigLIP 2 (NaFlex) | So400m/16 |   var. | [link](https://storage.googleapis.com/big_vision/siglip2/siglip2_so400m16_naflex.npz) |          83.5 |        55.1 |        71.2 |
+
+*The NaFlex results are for sequence length 256.*
+
+### Citation
+```
+@article{tschannen2025siglip,
+  title={SigLIP 2: Multilingual Vision-Language Encoders with Improved Semantic Understanding, Localization, and Dense Features},
+  author={Tschannen, Michael and Gritsenko, Alexey and Wang, Xiao and Naeem, Muhammad Ferjad and Alabdulmohsin, Ibrahim and Parthasarathy, Nikhil and Evans, Talfan and Beyer, Lucas and Xia, Ye and Mustafa, Basil and H\'enaff, Olivier and Harmsen, Jeremiah and Steiner, Andreas and Zhai, Xiaohua},
+  year={2025},
+  journal={arXiv preprint arXiv:2502.14786}
+}
+```
+
+\
+\
+Copyright 2025 Google LLC
+
+All software is licensed under the Apache License, Version 2.0 (Apache 2.0); you may not use this file except in compliance with the Apache 2.0 license. You may obtain a copy of the Apache 2.0 license at: https://www.apache.org/licenses/LICENSE-2.0
+
+All other materials are licensed under the Creative Commons Attribution 4.0 International License (CC-BY). You may obtain a copy of the CC-BY license at: https://creativecommons.org/licenses/by/4.0/legalcode
+
+Unless required by applicable law or agreed to in writing, all software and materials distributed here under the Apache 2.0 or CC-BY licenses are distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the licenses for the specific language governing permissions and limitations under those licenses.
+
+This is not an official Google product.

Tipsomaly/model/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,
+  })

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