Add files using upload-large-folder tool

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

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

Tipsomaly/model/big_vision/configs/proj/reward_tune/detection_reward.py

@@ -0,0 +1,232 @@
+# 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.
+
+"""Object detection reward from "Tuning computer vision models with task rewards" (https://arxiv.org/abs/2302.08242).
+
+The `reward_fn` computes the reward for a batch of predictions and ground truth
+annotations. When using it to optimize a model that outputs a prediction as a
+sequence of tokens like [y0, x0, Y0, X0, class0, confidence0, y1, x1, Y1, ...]
+the training loop may look like:
+
+```
+# Settings used in the paper.
+config.max_level = 1000  # Coordinates are discretized into 1000 buckets.
+config.max_conf = 2 # Two tokens are reserved to represent confidence.
+config.num_cls = 80 # Number of classes in COCO.
+config.nms_w = 0.3  # Weight for duplicate instances.
+config.cls_smooth = 0.05  # Adjust the classes weights based on their frequency.
+config.reward_thr = (0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95)
+config.correct_thr = 0.5  # Learn the IoU when matching with threshold=0.5.
+config.conf_w = 0.3  # Weight for the confidence loss.
+
+
+# 1) Sample N outputs for each input and compute rewards, use one sample to
+# optimize and others to compute a reward baseline.
+sample_seqs = sample_fn(params, images, num_samples)
+sample_rewards, aux = reward_fn(sample_seqs, labels, config)
+labels = sample_seqs[:, 0, ...]
+rewards = sample_rewards[:, 0]
+match_iou = aux["match_iou"][:, 0]
+baselines = (jnp.sum(sample_rewards, axis=-1) - rewards) / (num_samples - 1)
+
+# 2) Optimizize the model. By using REINFORCE to adjust the likelihood of the
+# sequence based on the reward and with supervision to teach the model to
+# predict the expected IoU of each box in its own samples.
+def loss_fn(params):
+  logits = model.apply(params, images, labels, train=True, rngs=rngs)
+  logits_softmax = jax.nn.log_softmax(logits)
+
+  # Use reinforce to optimize the expected reward for the whole sequence.
+  seq_rewards = (rewards - baselines)
+  # Note: consider improve this code to skip this loss for confidence tokens.
+  # The paper did not do it due to a bug (and also does not seem to matter).
+  target = jax.nn.one_hot(labels, logits.shape[-1]) * seq_rewards[:, None, None]
+  loss_reward = -jnp.sum(target * logits_softmax, axis=-1)
+
+  # Use supervision loss to tune the confidence tokens to predict IoU:
+  #   - (1.0, 0.0, 0.0, ...)   -> for padded boxes.
+  #   - (0.0, 1-iou, iou, ...) -> for sampled boxes.
+  conf0 = (labels[:, 5::6] == 0)
+  conf1 = (labels[:, 5::6] > 0) * (1.0 - match_iou)
+  conf2 = (labels[:, 5::6] > 0) * match_iou
+  target_conf = jnp.stack([conf0, conf1, conf2], axis=-1)
+  logits_conf = logits_softmax[:, 5::6, :3]
+  loss_conf = -jnp.sum(target_conf * logits_conf, axis=-1)
+
+  loss = jnp.mean(loss_reward) + config.conf_w * jnp.mean(loss_conf)
+  return loss
+```
+"""
+import functools
+
+import einops
+import jax
+import jax.numpy as jnp
+
+
+# Frequency of COCO object detection classes as observed in the training set.
+# pylint: disable=bad-whitespace,bad-continuation
+CLS_COUNTS = [
+    262465,  7113,  43867,   8725,   5135,   6069,   4571,   9973,  10759,
+     12884,  1865,   1983,   1285,   9838,  10806,   4768,   5508,   6587,
+      9509,  8147,   5513,   1294,   5303,   5131,   8720,  11431,  12354,
+      6496,  6192,   2682,   6646,   2685,   6347,   9076,   3276,   3747,
+      5543,  6126,   4812,  24342,   7913,  20650,   5479,   7770,   6165,
+     14358,  9458,   5851,   4373,   6399,   7308,   7852,   2918,   5821,
+      7179,  6353,  38491,   5779,   8652,   4192,  15714,   4157,   5805,
+      4970,  2262,   5703,   2855,   6434,   1673,   3334,    225,   5610,
+      2637, 24715,   6334,   6613,   1481,   4793,    198,   1954
+]
+# pylint: enable=bad-whitespace,bad-continuation
+
+
+def seq2box(seq, max_level, max_conf, num_cls):
+  """Extract boxes encoded as sequences."""
+  # Reshape to instances of boxes
+  dim_per_box = 6
+  seq_len = seq.shape[-1]
+  seq = seq[..., :(seq_len - seq_len % dim_per_box)]
+  seq = einops.rearrange(seq, "... (n d) -> ... n d", d=dim_per_box)
+
+  # Unpack box fields
+  boxes, labels, confs = seq[..., 0:4], seq[..., 4], seq[..., 5]
+  boxes = boxes - max_conf - 1
+  labels = labels - max_conf - 1 - max_level - 1
+  boxes = jnp.clip(boxes, 0, max_level) / max_level
+  labels = jnp.clip(labels, 0, num_cls - 1)
+  confs = jnp.clip(confs, 0, max_conf)
+
+  return boxes, labels, confs
+
+
+def iou_fn(box1, box2):
+  """Compute IoU of two boxes."""
+  ymin1, xmin1, ymax1, xmax1 = box1
+  ymin2, xmin2, ymax2, xmax2 = box2
+
+  a1 = jnp.abs((ymax1 - ymin1) * (xmax1 - xmin1))
+  a2 = jnp.abs((ymax2 - ymin2) * (xmax2 - xmin2))
+
+  yl = jnp.maximum(ymin1, ymin2)
+  yr = jnp.minimum(ymax1, ymax2)
+  yi = jnp.maximum(0, yr - yl)
+
+  xl = jnp.maximum(xmin1, xmin2)
+  xr = jnp.minimum(xmax1, xmax2)
+  xi = jnp.maximum(0, xr - xl)
+
+  inter = xi * yi
+  return inter / (a1 + a2 - inter + 1e-9)
+
+iou_fn_batched = jax.vmap(
+    jax.vmap(iou_fn, in_axes=(None, 0)), in_axes=(0, None)
+)
+
+
+def _reward_fn_thr(seq_pred, seq_gt,
+                   thr, nms_w, max_level, max_conf, num_cls, cls_smooth):
+  """Compute detection reward function for a given IoU threshold."""
+  # Weight matches of each label inversely proportional to the percentage of
+  # GT instances with such label in the whole train dataset. Additionally
+  # smooth out the observed distribution.
+  cls_counts = jnp.array(CLS_COUNTS)
+  weights = 1.0 / (cls_counts + cls_smooth*jnp.sum(cls_counts))
+  weights = num_cls * weights / jnp.sum(weights)
+
+  boxes_pred, labels_pred, confs_pred = seq2box(
+      seq_pred, max_level, max_conf, num_cls)
+  boxes_gt, labels_gt, confs_gt = seq2box(
+      seq_gt, max_level, max_conf, num_cls)
+
+  # Compute IoU matrix: Predictions X GT
+  iou = iou_fn_batched(boxes_pred, boxes_gt)
+
+  # IoU thr
+  iou = jnp.where(iou > thr, iou, 0.0)
+
+  # EOS mask
+  confs_mask = (confs_pred[:, None] > 0) * (confs_gt[None, :] > 0)
+  iou = confs_mask * iou
+
+  # Label mask
+  label_mask = labels_pred[:, None] == labels_gt[None, :]
+  iou = label_mask * iou
+
+  # Each prediction is matched to a single box
+  single_match_mask = jax.nn.one_hot(jnp.argmax(iou, axis=1), iou.shape[1])
+  iou = iou * single_match_mask
+
+  # Pred. boxes indicators
+  correct = jnp.any(iou > 0.0, axis=1).astype("int32") + 1
+  correct = jnp.where(confs_pred > 0, correct, 0)
+
+  # For each GT box find best match
+  matches_idx = jnp.argmax(iou, axis=0)
+  matches_iou = jnp.take_along_axis(iou, matches_idx[None], axis=0)[0]
+  matches_idx = jnp.where(matches_iou > 0.0, matches_idx, -1)
+
+  match_reward = jnp.sum((matches_idx >= 0) * weights[labels_gt][None, :])
+
+  # Compute duplicate penalty (aka NMS).
+  matches_mask = jax.nn.one_hot(matches_idx, iou.shape[0], axis=0)
+  nms_penalty = jnp.sum(
+      (iou > 0.0) * (1 - matches_mask) * weights[labels_pred][:, None])
+
+  match_iou = jnp.sum(iou, axis=1)
+
+  return {
+      "reward": (match_reward - nms_w * nms_penalty),
+      "num_matches": jnp.sum(matches_idx >= 0),
+      "nms_penalty": nms_penalty,
+      "correct": correct,
+      "match_iou": match_iou,
+  }
+
+
+def reward_fn(seqs_pred, seqs_gt, config):
+  """Total reward."""
+  result = {}
+  thrs = config.reward_thr
+  correct_thr = config.correct_thr
+  r_keys = ["reward", "num_matches", "nms_penalty"]
+  for thr in thrs:
+    fn = functools.partial(
+        _reward_fn_thr,
+        thr=thr,
+        nms_w=config.nms_w,
+        max_level=config.max_level,
+        max_conf=config.max_conf,
+        num_cls=config.num_cls,
+        cls_smooth=config.cls_smooth,
+    )
+    rewards = jax.vmap(jax.vmap(fn, in_axes=(0, None)))(seqs_pred, seqs_gt)
+
+    result = {**result, **{f"{k}-{thr:0.1f}": rewards[k]
+                           for k in r_keys}}
+    if thr == correct_thr:
+      correct = rewards["correct"]
+      match_iou = rewards["match_iou"]
+
+  result = {
+      **result,
+      **{k: jnp.mean(
+          jnp.array([result[f"{k}-{thr:0.1f}"] for thr in thrs]), axis=0)
+         for k in r_keys}
+  }
+
+  return result["reward"], {
+      "result": result,
+      "correct": correct,
+      "match_iou": match_iou,
+  }

Tipsomaly/model/big_vision/configs/proj/scaling_laws/train_vit_g.py

@@ -0,0 +1,87 @@
+# 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-train ViT-g (1B params) on JFT-3B as in https://arxiv.org/abs/2106.04560
+
+To train ViT-G (2B params), simply update the following single line:
+  `config.model.variant = 'G/14'`
+
+The code is released for reference purposes.
+One can test the code using public ImageNet-1k or ImageNet-21k dataset.
+
+big_vision.train \
+    --config big_vision/configs/proj/scaling_laws/train_vit_g.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():
+  """Rocket config."""
+  config = mlc.ConfigDict()
+
+  config.dataset = 'jft_3b'
+  config.val_split = 'val'
+  config.train_split = 'train'
+  config.num_classes = 29_593
+  config.init_head_bias = -10.0
+
+  # Fits 32 images per TPUv3 core with ViT-g/14.
+  config.batch_size = 4096*4
+
+  pp_common = '|value_range(-1, 1)'
+  pp_common += f'|onehot({config.num_classes})'
+  pp_common += '|keep("image", "labels")'
+  config.pp_train = 'inception_crop(224)|flip_lr' + pp_common
+  config.pp_eval = 'resize_small(256)|central_crop(224)' + pp_common
+  config.shuffle_buffer_size = 250_000  # Per host, so small-ish is ok.
+
+  config.log_training_steps = 50
+  config.log_eval_steps = 1000
+  # NOTE: eval is very fast O(seconds) so it's fine to run it often.
+
+  config.ckpt_steps = 1000
+  config.keep_ckpt_steps = 10_000
+
+  config.prefetch_to_device = 1
+  config.trial = 0
+
+  # Model section
+  config.model_name = 'vit'
+  config.model = mlc.ConfigDict()
+  config.model.variant = 'g/14'
+  config.model.pool_type = 'map'
+
+  # Optimizer section
+  config.optax_name = 'big_vision.scale_by_adafactor'
+  config.grad_clip_norm = 1.0
+  config.lr = 8e-4
+  config.wd = 0.03 * 8e-4
+  config.wd_mults = [
+      ('.*head/kernel', 100.0),
+      ('.*/kernel', 1.0),
+  ]
+  config.schedule = dict(
+      decay_type='rsqrt', timescale=10_000, warmup_steps=10_000,
+      cooldown_steps=50_000)
+  config.total_steps = 1_000_000
+
+  # Few-shot eval section
+  config.evals = {}
+  config.evals.fewshot = dict(log_steps=10_000, **get_fewshot_lsr())
+
+  return config

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

@@ -0,0 +1,84 @@
+# UViM: A Unified Modeling Approach for Vision with Learned Guiding Codes
+
+*by Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, Neil Houlsby*
+
+We provide pretrained UViM models from the [original paper](https://arxiv.org/abs/2205.10337),
+as well as the instructions on how to reproduce core paper experiments.
+
+## Pretrained models
+
+The table below contains UViM models (stage I and II) trained for three
+different tasks: panoptic segmentation, colorization and depth prediction.
+
+| task                  | model               | dataset                                                                  | accuracy     | download link                                                                             |
+| --------------------- | ------------------- | ------------------------------------------------------------------------ | ------------ | ----------------------------------------------------------------------------------------- |
+| Panoptic segmentation | UViM Stage I model  | [COCO(2017)](https://cocodataset.org/#home)                              |  75.8 PQ     | [link](https://storage.googleapis.com/big_vision/uvim/panoptic_stageI_params.npz)         |
+| Panoptic segmentation | UViM Stage II model | [COCO(2017)](https://cocodataset.org/#home)                              |  43.1 PQ     | [link](https://storage.googleapis.com/big_vision/uvim/panoptic_stageII_params.npz)        |
+| Colorization          | UViM Stage I model  | [ILSVRC-2012](https://www.image-net.org/)                                |  15.59 FID   | [link](https://storage.googleapis.com/big_vision/uvim/color_stageI_params.npz)            |
+| Colorization          | UViM Stage II model | [ILSVRC-2012](https://www.image-net.org/)                                |  16.99 FID   | [link](https://storage.googleapis.com/big_vision/uvim/color_stageII_params.npz)           |
+| Depth                 | UViM Stage I model  | [NYU Depth V2](https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html) |  0.155 RMSE  | [link](https://storage.googleapis.com/big_vision/uvim/depth_stageI_params.npz)            |
+| Depth                 | UViM Stage II model | [NYU Depth V2](https://cs.nyu.edu/~silberman/datasets/nyu_depth_v2.html) |  0.463 RMSE  | [link](https://storage.googleapis.com/big_vision/uvim/depth_stageII_params.npz)           |
+
+All of this models can be interactively explored in our [colabs](configs/proj/uvim).
+
+## Running on a single-host TPU machine
+
+Below we provide instructions on how to run UViM training (stage I and
+stage II) using a single TPU host with 8 TPU accelerators. These instructions
+can be easily adapted to a GPU host and multi-host TPU setup, see the main
+`big_vision` [README file](README.md).
+
+We assume that the user has already created and `ssh`-ed to the TPU host
+machine. The next step is to clone `big_vision` repository:
+`git clone https://github.com/google-research/big_vision.git`.
+
+The next steps are to create a python virtual environment and install python
+dependencies:
+```
+virtualenv bv
+source bv/bin/activate
+cd big_vision/
+pip3 install --upgrade pip
+pip3 install -r big_vision/requirements.txt
+pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+```
+
+After this invoke the helper tool to download and prepare data:
+`python3 -m big_vision.tools.download_tfds_datasets coco/2017_panoptic nyu_depth_v2`.
+For preparing the ImageNet dataset consult the main codebase README.
+
+> :warning: TPU machines have 100 GB of the disk space. It may not be enough to
+> store all training data (though only panoptic or only depth data may fit).
+> Consider preparing the data on a seperate machine and then copying it to
+> to TPU machine's extra persistent disk or to a Google Cloud Bucket. See
+> instructions for [creating an extra persistent disk](https://cloud.google.com/tpu/docs/users-guide-tpu-vm).
+> Remember to set the correct data home directory, e.g.`export DISK=/mnt/disk/persist; export TFDS_DATA_DIR=$DISK/tensorflow_datasets`.
+
+Our panoptic evaluator uses raw variant of the COCO data, so we move it into a
+separate folder. Note, `tfds` has already pre-downloaded the panoptic data,
+except for one small json file that we fetch manually:
+```
+mkdir $DISK/coco_data
+cd $DISK/coco_data
+mv $TFDS_DATA_DIR/downloads/extracted/ZIP.image.cocod.org_annot_panop_annot_train<REPLACE_ME_WITH_THE_HASH_CODE>.zip/annotations/* .
+wget https://raw.githubusercontent.com/cocodataset/panopticapi/master/panoptic_coco_categories.json
+export COCO_DATA_DIR=$DISK/coco_data
+```
+
+For FID evaluator, which is used for the colorization model, set the path to the
+directory with image id files, e.g.
+`export FID_DATA_DIR=<ROOT>/big_vision/evaluators/proj/uvim/coltran_fid_data`.
+
+As an example, stage I panoptic training can be invoked as (note the `:singlehost` config parameter which will use lightweight configuration suitable for a single host):
+```
+python3 -m big_vision.trainers.proj.uvim.vqvae --config big_vision/configs/proj/uvim/vqvae_coco_panoptic.py:singlehost --workdir workdirs/`date '+%m-%d_%H%M'`
+```
+or stage II training
+```
+python3 -m big_vision.trainers.proj.uvim.train --config big_vision/configs/proj/uvim/train_coco_panoptic_pretrained.py:singlehost --workdir workdirs/`date '+%m-%d_%H%M'`
+```
+
+## Acknowledgments
+The sampling code in `models/proj/uvim/decode.py` module is based on contributions
+from Anselm Levskaya, Ilya Tolstikhin and Maxim Neumann.
+

Tipsomaly/model/big_vision/configs/proj/uvim/train_coco_panoptic_pretrained.py

@@ -0,0 +1,164 @@
+# 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"""A config for training a UViM stage II model for the panoptic task.
+
+This config is expected to reproduce the paper's result and achieve
+approximately 43.7 PQ points on the COCO holdout data.
+
+We also provide a low-resource variant of this config, which can be enabled
+by adding `:singlehost` postfix to the config name. This one is expected to
+achieve 39.4 PQ points on the COCO holdout data.
+"""
+
+import big_vision.configs.common as bvcc
+from ml_collections import ConfigDict
+
+VTT_MODELS = {
+    'base': dict(num_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768),
+    'large': dict(num_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024),
+}
+
+VQVAE_MODELS = {
+    'base': dict(enc_depth=6, dec_depth=12, num_heads=12, mlp_dim=3072, width=768),
+}
+
+RES = 512
+PATCH_SIZE = 16
+LABEL_RES = 512
+LABEL_PATCH_SIZE = 16
+
+
+def get_config(arg=''):
+  """Config for training."""
+  arg = bvcc.parse_arg(arg, runlocal=False, singlehost=False)
+  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({LABEL_RES}, inkey="image", outkey="image_ctx")|'
+      f'resize({RES})|resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1, 1, key="image_ctx")|'
+      f'value_range(-1, 1)|make_canonical|keep("image","image_ctx","labels")'
+  )
+  pp_eval = (
+      f'decode|coco_panoptic|concat(["semantics","instances"], "labels")|'
+      f'resize({LABEL_RES}, inkey="image", outkey="image_ctx")|'
+      f'resize({RES})|resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1, 1, key="image_ctx")|'
+      f'value_range(-1, 1)|make_canonical|keep("image","image_ctx","labels")'
+  )
+  pp_predict = (
+      f'resize({LABEL_RES}, inkey="image", outkey="image_ctx")|resize({RES})|'
+      f'value_range(-1, 1, key="image_ctx")|value_range(-1, 1)|'
+      f'keep("image","image_ctx","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 = 5000
+  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.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.oracle = ConfigDict()
+  config.oracle.task = 'proj.uvim.panoptic_task'
+  config.oracle.model_init = 'gs://big_vision/uvim/panoptic_stageI_params.npz'
+  config.oracle.model_name = 'proj.uvim.vit'
+  config.oracle.model = ConfigDict(VQVAE_MODELS['base'])
+  config.oracle.model.input_size = (LABEL_RES, LABEL_RES)
+  config.oracle.model.patch_size = (LABEL_PATCH_SIZE, LABEL_PATCH_SIZE)
+  config.oracle.model.code_len = 256
+  config.oracle.model.dict_size = 4096
+  config.oracle.model.codeword_dim = 768
+  config.oracle.model.with_encoder_ctx = True
+  config.oracle.model.with_decoder_ctx = True
+  config.oracle.model.code_dropout = 'random'
+  config.oracle.model.bottleneck_resize = True
+  config.oracle.model.inputs = {
+      'semantics': (133 + 1, LABEL_PATCH_SIZE**2),  # +1 for void label
+      'instances': (100, LABEL_PATCH_SIZE**2),  # COCO: actually 98 train/78 validation.
+  }
+  config.oracle.model.outputs = config.oracle.model.inputs
+
+  # Model section
+  config.model_name = 'proj.uvim.vtt'
+  # config.model_init = {'encoder': 'howto-i21k-B/8'}
+  config.model_init = {'encoder': 'howto-i21k-L/16'}
+  config.model = ConfigDict(VTT_MODELS['large'])
+  config.model.patches = ConfigDict({'size': (PATCH_SIZE, PATCH_SIZE)})
+  config.model.vocab_size = config.oracle.model.get_ref('dict_size') + 1
+  config.model.posemb_type = 'learn'
+  config.model.input_size = (RES, RES)
+  config.model.seq_len = config.oracle.model.get_ref('code_len')
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = ConfigDict()
+  config.evals.val.type = 'proj.uvim.compute_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
+
+  base = {
+      'type': 'proj.uvim.coco_panoptic',
+      'pp_fn': pp_predict,
+      'log_steps': 10_000,
+      # Filters objects that occupy less than 0.03^2 fraction of all pixels.
+      # 'predict_kwargs': {'min_fraction': 0.03 ** 2},
+  }
+  config.evals.coco_panoptic_train = dict(**base, split='train[4096:8192]')
+  config.evals.coco_panoptic_holdout = dict(**base, split='train[:4096]')
+  config.evals.coco_panoptic = dict(**base, split='validation')
+
+  # config.evals.save_pred = dict(type='proj.uvim.save_predictions')
+  # config.evals.save_pred.pp = pp_eval.replace('decode|', '')
+  # config.evals.save_pred.log_steps = 100_000
+  # config.evals.save_pred.dataset = config.dataset
+  # config.evals.save_pred.split = 'validation[:1024]'
+  # config.evals.save_pred.outfile = 'inference.npz'
+
+  if arg.singlehost:
+    config.input.batch_size = 32
+    config.num_epochs = 50
+  elif arg.runlocal:
+    config.input.batch_size = 4
+    config.input.shuffle_buffer_size = 10
+    config.evals.val.data.split = 'train[:16]'
+  return config

Tipsomaly/model/big_vision/configs/proj/uvim/train_imagenet2012_colorization_pretrained.py

@@ -0,0 +1,161 @@
+# 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"""A config for training a UViM stage II model for the colorization task.
+"""
+
+import big_vision.configs.common as bvcc
+from ml_collections import ConfigDict
+
+VTT_MODELS = {
+    'base': dict(num_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768),
+    'large': dict(num_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024),
+}
+
+VQVAE_MODELS = {
+    'base': dict(enc_depth=6, dec_depth=12, num_heads=12, mlp_dim=3072, width=768),
+}
+
+RES = 512
+PATCH_SIZE = 16
+LABEL_RES = 512
+LABEL_PATCH_SIZE = 16
+
+
+def get_config(arg=''):
+  """Config for training."""
+  arg = bvcc.parse_arg(arg, runlocal=False, singlehost=False)
+  config = ConfigDict()
+
+  config.input = {}
+  config.input.pp = (
+      f'decode_jpeg_and_inception_crop({RES})'
+      f'|flip_lr'
+      f'|copy(inkey="image", outkey="labels")'
+      f'|resize({LABEL_RES},inkey="labels",outkey="labels",method="nearest")'
+      f'|value_range(-1,1,key="labels")'
+      f'|rgb_to_grayscale_to_rgb(inkey="image",outkey="image")'
+      f'|value_range(-1,1,key="image")'
+      f'|copy(inkey="image", outkey="image_ctx")'
+      f'|resize({LABEL_RES},inkey="image_ctx",outkey="image_ctx")'
+      f'|keep("image","image_ctx","labels")')
+  pp_eval = (
+      f'decode'
+      f'|resize({RES})'
+      f'|copy(inkey="image", outkey="labels")'
+      f'|resize({LABEL_RES},inkey="labels",outkey="labels",method="nearest")'
+      f'|value_range(-1,1,key="labels")'
+      f'|rgb_to_grayscale_to_rgb(inkey="image",outkey="image")'
+      f'|value_range(-1,1,key="image")'
+      f'|copy(inkey="image", outkey="image_ctx")'
+      f'|resize({LABEL_RES},inkey="image_ctx",outkey="image_ctx")'
+      f'|strong_hash(inkey="tfds_id", outkey="image/id")'
+      f'|keep("image","image_ctx","labels","image/id")')
+
+  config.input.data = dict(name='imagenet2012', split='train[4096:]')
+  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 = 5000
+  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.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.oracle = ConfigDict()
+  config.oracle.task = 'proj.uvim.colorization_task'
+  config.oracle.model_init = 'gs://big_vision/uvim/color_stageI_params.npz'
+  config.oracle.model_name = 'proj.uvim.vit'
+  config.oracle.model = ConfigDict(VQVAE_MODELS['base'])
+  config.oracle.model.input_size = (LABEL_RES, LABEL_RES)
+  config.oracle.model.patch_size = (LABEL_PATCH_SIZE, LABEL_PATCH_SIZE)
+  config.oracle.model.code_len = 256
+  config.oracle.model.dict_size = 4096
+  config.oracle.model.codeword_dim = 768
+  config.oracle.model.with_encoder_ctx = True
+  config.oracle.model.with_decoder_ctx = True
+  config.oracle.model.code_dropout = 'random'
+  config.oracle.model.bottleneck_resize = True
+  config.oracle.model.inputs = {
+      'color': (3, LABEL_PATCH_SIZE**2),
+  }
+  config.oracle.model.outputs = config.oracle.model.inputs
+
+  # Model section
+  config.model_name = 'proj.uvim.vtt'
+  # config.model_init = {'encoder': 'howto-i21k-B/8'}
+  config.model_init = {'encoder': 'howto-i21k-L/16'}
+  config.model = ConfigDict(VTT_MODELS['large'])
+  config.model.patches = ConfigDict({'size': (PATCH_SIZE, PATCH_SIZE)})
+  config.model.vocab_size = config.oracle.model.get_ref('dict_size') + 1
+  config.model.posemb_type = 'learn'
+  config.model.input_size = (RES, RES)
+  config.model.seq_len = config.oracle.model.get_ref('code_len')
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = ConfigDict()
+  config.evals.val.type = 'proj.uvim.compute_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
+
+  base = {
+      'type': 'proj.uvim.psnr',
+      'pp_fn': pp_eval.replace('decode|', ''),
+      'log_steps': 10_000,
+  }
+  config.evals.psnr_train = dict(**base, split='train[4096:8192]')
+  config.evals.psnr_holdout = dict(**base, split='train[:4096]')
+  config.evals.psnr_val = dict(**base, split='validation')
+
+  config.evals.colorization_val_coltran_fid = {
+      'type': 'proj.uvim.coltran_fid',
+      'log_steps': 100_000,
+  }
+
+  # config.evals.save_pred = dict(type='proj.uvim.save_predictions')
+  # config.evals.save_pred.pp_fn = pp_eval.replace('decode|', '')
+  # config.evals.save_pred.log_steps = 100_000
+  # config.evals.save_pred.dataset = config.dataset
+  # config.evals.save_pred.split = 'validation[:1024]'
+  # config.evals.save_pred.outfile = 'inference.npz'
+
+  if arg.singlehost:
+    config.input.batch_size = 32
+    config.total_epochs = 20
+  elif arg.runlocal:
+    config.input.batch_size = 8
+    config.input.shuffle_buffer_size = 10
+    config.evals.val.data.split = 'validation[:256]'
+  return config

Tipsomaly/model/big_vision/configs/proj/uvim/train_nyu_depth_pretrained.py

@@ -0,0 +1,170 @@
+# 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"""A config for training a UViM stage II model for the depth task.
+"""
+
+import big_vision.configs.common as bvcc
+from ml_collections import ConfigDict
+
+
+VTT_MODELS = {
+    'base': dict(num_layers=12, num_heads=12, mlp_dim=3072, emb_dim=768),
+    'large': dict(num_layers=24, num_heads=16, mlp_dim=4096, emb_dim=1024),
+}
+
+VQVAE_MODELS = {
+    'base': dict(enc_depth=6, dec_depth=12, num_heads=12, mlp_dim=3072, width=768),
+}
+
+
+RES = 512
+PATCH_SIZE = 16
+LABEL_RES = 512
+LABEL_PATCH_SIZE = 16
+QUANTIZATION_BINS = 256
+# Same as values used in eval, see evaluators/nyu_depth.py.
+MIN_DEPTH = 1e-3
+MAX_DEPTH = 10
+
+
+def get_config(arg='split=final'):
+  """Config for training."""
+  arg = bvcc.parse_arg(arg, split='final', 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},inkey="image",outkey="image_ctx")|'
+      f'resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1,1)|'
+      f'value_range(-1,1,inkey="image_ctx",outkey="image_ctx")|'
+      f'keep("image","image_ctx","labels")'
+  )
+  pp_eval = (
+      f'decode|nyu_depth|'
+      f'nyu_eval_crop|'
+      f'resize({RES})|'
+      f'resize({LABEL_RES},inkey="image",outkey="image_ctx")|'
+      f'resize({LABEL_RES},key="labels",method="nearest")|'
+      f'value_range(-1,1)|'
+      f'value_range(-1,1,inkey="image_ctx",outkey="image_ctx")|'
+      f'keep("image","image_ctx","labels")'
+  )
+  pp_predict = (
+      f'nyu_depth|'
+      f'nyu_eval_crop|copy("labels","ground_truth")|'
+      f'resize({RES})|'
+      f'resize({LABEL_RES},inkey="image",outkey="image_ctx")|'
+      f'value_range(-1,1)|'
+      f'value_range(-1,1,inkey="image_ctx",outkey="image_ctx")|'
+      f'keep("image","image_ctx","ground_truth")'
+  )
+
+  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 = 5000
+  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.optax.clipping_threshold = None
+
+  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.oracle = ConfigDict()
+  config.oracle.min_depth = MIN_DEPTH
+  config.oracle.max_depth = MAX_DEPTH
+  config.oracle.task = 'proj.uvim.depth_task'
+  config.oracle.model_init = 'gs://big_vision/uvim/depth_stageI_params.npz'
+  config.oracle.model_name = 'proj.uvim.vit'
+  config.oracle.model = ConfigDict(VQVAE_MODELS['base'])
+  config.oracle.model.input_size = (LABEL_RES, LABEL_RES)
+  config.oracle.model.patch_size = (LABEL_PATCH_SIZE, LABEL_PATCH_SIZE)
+  config.oracle.model.code_len = 256
+  config.oracle.model.dict_size = 4096
+  config.oracle.model.codeword_dim = 768
+  config.oracle.model.with_encoder_ctx = True
+  config.oracle.model.with_decoder_ctx = True
+  config.oracle.model.code_dropout = 'random'
+  config.oracle.model.bottleneck_resize = True
+  config.oracle.model.inputs = {
+      'depth': (QUANTIZATION_BINS, LABEL_PATCH_SIZE**2,),
+  }
+  config.oracle.model.outputs = config.oracle.model.inputs
+
+  # Model section
+  config.model_name = 'proj.uvim.vtt'
+  # config.model_init = {'encoder': 'howto-i21k-B/8''}  # B/8 I21K
+  config.model_init = {'encoder': 'howto-i21k-L/16'}  # L/16 I21K
+  config.model = ConfigDict(VTT_MODELS['large'])
+  config.model.patches = ConfigDict({'size': (PATCH_SIZE, PATCH_SIZE)})
+  config.model.vocab_size = config.oracle.model.dict_size + 1
+  config.model.posemb_type = 'learn'
+  config.model.input_size = (RES, RES)
+  config.model.seq_len = config.oracle.model.get_ref('code_len')
+  config.model.zero_decoder_seq = False
+
+  # Evaluation section
+  config.evals = {}
+  config.evals.val = ConfigDict()
+  config.evals.val.type = 'proj.uvim.compute_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 = 1000
+
+  base = {
+      'type': 'proj.uvim.nyu_depth',
+      'dataset': config.input.data.name,
+      'pp_fn': pp_predict,
+      'log_steps': 2000,
+      'min_depth': MIN_DEPTH,
+      'max_depth': MAX_DEPTH,
+  }
+  config.evals.nyu_depth_val = dict(**base, split='validation')
+
+  if arg.singlehost:
+    config.input.batch_size = 32
+    config.total_epochs = 20
+  elif arg.runlocal:
+    config.oracle.model_init = '/tmp/checkpoint.npz'
+    config.model_init = {'encoder': '/tmp/enc_checkpoint.npz'}
+    config.evals = {}
+    config.input.batch_size = 1
+    config.input.shuffle_buffer_size = 10
+  return config

Tipsomaly/model/big_vision/configs/proj/uvim/uvim_color_task.ipynb

@@ -0,0 +1,167 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "name": "UViM color task",
+      "provenance": [],
+      "collapsed_sections": [],
+      "private_outputs": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU",
+    "gpuClass": "standard"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "source": [
+        "# Fetch big_vision repository and move it into the current workdir (import path).\n",
+        "!git clone --depth=1 https://github.com/google-research/big_vision big_vision_repo\n",
+        "!cp -R big_vision_repo/big_vision big_vision\n",
+        "!pip install -qr big_vision/requirements.txt"
+      ],
+      "metadata": {
+        "id": "sKZK6_QpVI_O"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "import numpy as np\n",
+        "\n",
+        "from big_vision.models.proj.uvim import vtt  # stage-II model\n",
+        "from big_vision.models.proj.uvim import vit  # stage-I model\n",
+        "\n",
+        "from big_vision.models.proj.uvim import decode\n",
+        "from big_vision.trainers.proj.uvim import colorization_task as task\n",
+        "from big_vision.configs.proj.uvim import train_imagenet2012_colorization_pretrained as config_module\n",
+        "\n",
+        "import big_vision.pp.ops_image\n",
+        "import big_vision.pp.ops_general\n",
+        "import big_vision.pp.proj.uvim.pp_ops\n",
+        "from big_vision.pp import builder as pp_builder\n",
+        "\n",
+        "config = config_module.get_config()\n",
+        "res = 512\n",
+        "seq_len = config.model.seq_len\n",
+        "\n",
+        "lm_model = vtt.Model(**config.model)\n",
+        "oracle_model = vit.Model(**config.oracle.model)\n",
+        "\n",
+        "preprocess_fn = pp_builder.get_preprocess_fn(\n",
+        "    'decode|resize(512)|'\n",
+        "    'rgb_to_grayscale_to_rgb|value_range(-1,1)|'\n",
+        "    'copy(inkey=\"image\",outkey=\"image_ctx\")')\n",
+        "\n",
+        "@jax.jit\n",
+        "def predict_code(params, x, rng, temperature):\n",
+        "  prompts = jnp.zeros((x[\"image\"].shape[0], seq_len), dtype=jnp.int32)\n",
+        "  seqs, _, _ = decode.temperature_sampling(\n",
+        "      params=params, model=lm_model, seed=rng,\n",
+        "      inputs=x[\"image\"],\n",
+        "      prompts=prompts,\n",
+        "      temperature=temperature,\n",
+        "      num_samples=1, eos_token=-1, prefill=False)\n",
+        "  seqs = jnp.squeeze(seqs, axis=1)  # drop num_samples axis \n",
+        "  return seqs - 1\n",
+        "  \n",
+        "@jax.jit\n",
+        "def labels2code(params, x, ctx):\n",
+        "  y, aux = oracle_model.apply(params, x, ctx=ctx, train=False, method=oracle_model.encode)\n",
+        "  return aux[\"code\"]\n",
+        "\n",
+        "@jax.jit\n",
+        "def code2labels(params, code, ctx):\n",
+        "  logits, aux = oracle_model.apply(params, code, ctx=ctx, train=False, discrete_input=True, method=oracle_model.decode)\n",
+        "  return task.predict_outputs(logits, config.oracle)"
+      ],
+      "metadata": {
+        "id": "QzThueWDzc7I"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Load checkpoints\n",
+        "!gsutil cp -n gs://big_vision/uvim/color_stageI_params.npz gs://big_vision/uvim/color_stageII_params.npz .\n",
+        "\n",
+        "oracle_params, oracle_state = vit.load(None, \"color_stageI_params.npz\")\n",
+        "oracle_params = jax.device_put({\"params\": oracle_params, \"state\": oracle_state})\n",
+        "\n",
+        "lm_params = vtt.load(None, \"color_stageII_params.npz\")\n",
+        "lm_params = jax.device_put({\"params\": lm_params})"
+      ],
+      "metadata": {
+        "id": "AEjRgshLa6Fp"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Prepare set of images from coco/val2017:\n",
+        "#  - https://cocodataset.org/\n",
+        "import os\n",
+        "import tensorflow as tf\n",
+        "\n",
+        "if not os.path.exists(\"val2017/\"):\n",
+        "  !wget --no-clobber http://images.cocodataset.org/zips/val2017.zip\n",
+        "  !unzip -uq val2017.zip\n",
+        "\n",
+        "dataset = tf.data.Dataset.list_files(\"val2017/*.jpg\", shuffle=True)\n",
+        "dataset = dataset.map(lambda filename: {\"image\": tf.io.read_file(filename)})\n",
+        "dataset = dataset.map(preprocess_fn)"
+      ],
+      "metadata": {
+        "id": "BKifDDRnH_Ll"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# Run the model in a few examples:\n",
+        "from matplotlib import pyplot as plt\n",
+        "\n",
+        "num_examples = 4\n",
+        "data = dataset.batch(1).take(num_examples).as_numpy_iterator()\n",
+        "key = jax.random.PRNGKey(0)\n",
+        "temperature = jnp.array(1.0)\n",
+        "\n",
+        "def render_example(image, prediction):\n",
+        "  f, ax = plt.subplots(1, 2, figsize=(10, 10))\n",
+        "  ax[0].imshow(image*0.5 + 0.5)\n",
+        "  ax[0].axis(\"off\")\n",
+        "  ax[1].imshow(prediction*0.5 + 0.5)\n",
+        "  ax[1].axis(\"off\")\n",
+        "\n",
+        "for idx, batch in enumerate(data):\n",
+        "  subkey = jax.random.fold_in(key, idx)\n",
+        "  code = predict_code(lm_params, batch, key, temperature)\n",
+        "  aux_inputs = task.input_pp(batch, config.oracle)\n",
+        "  prediction = code2labels(oracle_params, code, aux_inputs[\"ctx\"])\n",
+        "  render_example(batch[\"image\"][0], prediction[\"color\"][0])"
+      ],
+      "metadata": {
+        "id": "TuevCy33nuv3"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}