Reset again!

.gitattributes

@@ -1,34 +0,0 @@
-*.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
-*.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

.gitignore

@@ -1,163 +0,0 @@
-cache/
-.idea
-
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-#   in version control.
-#   https://pdm.fming.dev/#use-with-ide
-.pdm.toml
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "prismer"]
-	path = prismer
-	url = https://github.com/nvlabs/prismer

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/prismer_demo.iml" filepath="$PROJECT_DIR$/.idea/prismer_demo.iml" />
+    </modules>
+  </component>
+</project>

.idea/prismer_demo.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/vcs.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+    <mapping directory="$PROJECT_DIR$/prismer" vcs="Git" />
+  </component>
+</project>

.pre-commit-config.yaml

@@ -1,36 +0,0 @@
-repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.2.0
-  hooks:
-  - id: check-executables-have-shebangs
-  - id: check-json
-  - id: check-merge-conflict
-  - id: check-shebang-scripts-are-executable
-  - id: check-toml
-  - id: check-yaml
-  - id: double-quote-string-fixer
-  - id: end-of-file-fixer
-  - id: mixed-line-ending
-    args: ['--fix=lf']
-  - id: requirements-txt-fixer
-  - id: trailing-whitespace
-- repo: https://github.com/myint/docformatter
-  rev: v1.4
-  hooks:
-  - id: docformatter
-    args: ['--in-place']
-- repo: https://github.com/pycqa/isort
-  rev: 5.12.0
-  hooks:
-    - id: isort
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.991
-  hooks:
-    - id: mypy
-      args: ['--ignore-missing-imports']
-      additional_dependencies: ['types-python-slugify']
-- repo: https://github.com/google/yapf
-  rev: v0.32.0
-  hooks:
-  - id: yapf
-    args: ['--parallel', '--in-place']

.style.yapf

@@ -1,5 +0,0 @@
-[style]
-based_on_style = pep8
-blank_line_before_nested_class_or_def = false
-spaces_before_comment = 2
-split_before_logical_operator = true

app.py

@@ -5,11 +5,20 @@ from __future__ import annotations
 import os
 import shutil
 import subprocess
+
 import gradio as gr
 
+if os.getenv('SYSTEM') == 'spaces':
+    with open('patch') as f:
+        subprocess.run('patch -p1'.split(), cwd='prismer', stdin=f)
+    shutil.copytree('prismer/helpers/images',
+                    'prismer/images',
+                    dirs_exist_ok=True)
+
 from app_caption import create_demo as create_demo_caption
 from prismer_model import build_deformable_conv, download_models
 
+
 # Prepare model checkpoints
 download_models()
 build_deformable_conv()

app_caption.py

@@ -15,8 +15,10 @@ def create_demo():
 
     with gr.Row():
         with gr.Column():
-            image = gr.Image(label='Input Image', type='filepath')
-            model_name = gr.Dropdown(label='Model Size', choices=['prismer_base'], value='prismer_base')
+            image = gr.Image(label='Input', type='filepath')
+            model_name = gr.Dropdown(label='Model',
+                                     choices=['prismer_base'],
+                                     value='prismer_base')
             run_button = gr.Button('Run')
         with gr.Column(scale=1.5):
             caption = gr.Text(label='Caption')
@@ -30,7 +32,15 @@ def create_demo():
                 ocr = gr.Image(label='OCR Detection')
 
     inputs = [image, model_name]
-    outputs = [caption, depth, edge, normals, segmentation, object_detection, ocr]
+    outputs = [
+        caption,
+        depth,
+        edge,
+        normals,
+        segmentation,
+        object_detection,
+        ocr,
+    ]
 
     paths = sorted(pathlib.Path('prismer/images').glob('*'))
     examples = [[path.as_posix(), 'prismer_base'] for path in paths]

patch

@@ -0,0 +1,82 @@
+diff --git a/dataset/caption_dataset.py b/dataset/caption_dataset.py
+index 266fdda..0cc5d3f 100644
+--- a/dataset/caption_dataset.py
++++ b/dataset/caption_dataset.py
+@@ -50,7 +50,7 @@ class Caption(Dataset):
+         elif self.dataset == 'demo':
+             img_path_split = self.data_list[index]['image'].split('/')
+             img_name = img_path_split[-2] + '/' + img_path_split[-1]
+-            image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts)
++            image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts)
+ 
+         experts = self.transform(image, labels)
+         experts = post_label_process(experts, labels_info)
+diff --git a/dataset/utils.py b/dataset/utils.py
+index b368aac..418358c 100644
+--- a/dataset/utils.py
++++ b/dataset/utils.py
+@@ -5,6 +5,7 @@
+ # https://github.com/NVlabs/prismer/blob/main/LICENSE
+ 
+ import os
++import pathlib
+ import re
+ import json
+ import torch
+@@ -14,10 +15,12 @@ import torchvision.transforms as transforms
+ import torchvision.transforms.functional as transforms_f
+ from dataset.randaugment import RandAugment
+ 
+-COCO_FEATURES = torch.load('dataset/coco_features.pt')['features']
+-ADE_FEATURES = torch.load('dataset/ade_features.pt')['features']
+-DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features']
+-BACKGROUND_FEATURES = torch.load('dataset/background_features.pt')
++cur_dir = pathlib.Path(__file__).parent
++
++COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
++ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
++DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
++BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
+ 
+ 
+ class Transform:
+diff --git a/model/prismer.py b/model/prismer.py
+index 080253a..02362a4 100644
+--- a/model/prismer.py
++++ b/model/prismer.py
+@@ -5,6 +5,7 @@
+ # https://github.com/NVlabs/prismer/blob/main/LICENSE
+ 
+ import json
++import pathlib
+ import torch.nn as nn
+ 
+ from model.modules.vit import load_encoder
+@@ -12,6 +13,9 @@ from model.modules.roberta import load_decoder
+ from transformers import RobertaTokenizer, RobertaConfig
+ 
+ 
++cur_dir = pathlib.Path(__file__).parent
++
++
+ class Prismer(nn.Module):
+     def __init__(self, config):
+         super().__init__()
+@@ -26,7 +30,7 @@ class Prismer(nn.Module):
+             elif exp in ['obj_detection', 'ocr_detection']:
+                 self.experts[exp] = 64
+ 
+-        prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']]
++        prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']]
+         roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model'])
+ 
+         self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name'])
+@@ -35,7 +39,7 @@ class Prismer(nn.Module):
+ 
+         self.prepare_to_train(config['freeze'])
+         self.ignored_modules = self.get_ignored_modules(config['freeze'])
+-    
++
+     def prepare_to_train(self, mode='none'):
+         for name, params in self.named_parameters():
+             if mode == 'freeze_lang':

prismer/.gitignore

@@ -0,0 +1,10 @@
+.idea
+cache
+.DS_Store
+**/__pycache__/*
+helpers/data/*
+helpers/images2/*
+helpers/labels/*
+experts/expert_weights
+logging/*
+flagged/*

prismer/LICENSE

@@ -0,0 +1,97 @@
+Copyright (c) 2023, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+NVIDIA Source Code License for Prismer
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor’s or its affiliates’ names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

prismer/README.md

@@ -0,0 +1,156 @@
+# Prismer
+
+This repository contains the source code of **Prismer** and **PrismerZ** from the paper, [Prismer: A Vision-Language Model with An Ensemble of Experts](https://arxiv.org/abs/2303.02506).
+
+<img src="helpers/intro.png"  width="100%"/>
+
+## Get Started
+The implementation is based on `PyTorch 1.13`, and highly integrated with Huggingface [`accelerate`](https://github.com/huggingface/accelerate) toolkit for readable and optimised multi-node multi-gpu training.
+
+First, let's install all package dependencies by running
+```bash
+pip install -r requirements.txt
+```
+
+### Prepare Accelerator Config
+Then we generate the corresponding `accelerate`  config based on your training server configuration. For both single-node multi-gpu and multi-node multi-gpu training, simply run
+```bash
+# to get your machine rank 0 IP address
+hostname -i
+
+# and for each machine, run the following command, set --num_machines 1 in a single-node setting
+python generate_config.py —-main_ip {MAIN_IP} -—rank {MACHINE_RANK} —-num_machines {TOTAL_MACHINES}
+```
+
+## Datasets
+
+### Pre-training
+We pre-train Prismer/PrismerZ with a combination of five widely used image-alt/text datasets, with pre-organised data lists provided below.
+- [COCO 2014](https://www.dropbox.com/s/6btr8hz5n1e1q4d/coco_karpathy_train.json?dl=0): the Karpathy training split (which will also be used for fine-tuning).
+- [Visual Genome](https://www.dropbox.com/s/kailbaay0sqraxc/vg_caption.json?dl=0): the official Visual Genome captioning dataset.
+- [CC3M + SGU](https://www.dropbox.com/s/xp2nuhc88f1czxm/filtered_cc3m_sbu.json?dl=0): filtered and re-captioned by BLIP-Large.
+- [CC12M](https://www.dropbox.com/s/th358bb6wqkpwbz/filtered_cc12m.json?dl=0): filtered and re-captioned by BLIP-Large.
+
+The web datasets (CC3M, SGU, CC12M) is composed with image urls. It is highly recommended to use [img2dataset](https://github.com/rom1504/img2dataset), a highly optimised toolkit for large-scale web scraping to download these images. An example bash script of using `img2dataset` to download `cc12m` dataset is provided below.
+```bash
+img2dataset --url_list filtered_cc12m.json --input_format "json" --url_col "url" --caption_col "caption" --output_folder cc12m --processes_count 16 --thread_count 64 --image_size 256
+```
+
+*Note: It is expected that the number of downloaded images is less than the number of images in the json file, because some urls might not be valid or require long loading time.*
+
+### Image Captioning / VQA
+We evaluate image captioning performance on two datasets, COCO 2014 and NoCaps; and VQA performance on VQAv2 dataset. In VQA tasks, we additionally augment the training data with Visual Genome QA, following BLIP. Again, we have prepared and organised the training and evaluation data lists provided below.
+
+- [Image Captioning](https://www.dropbox.com/sh/quu6v5hzdetjcdz/AACze0_h6BO8LJmSsEq4MM8-a?dl=0): including COCO (Karpathy Split) and NoCaps.
+- [VQAv2](https://www.dropbox.com/sh/hqtxl1k8gkbhhoi/AACiax5qi7no3pJgO1E57Xefa?dl=0): including VQAv2 and VG QA.
+
+## Generating Expert Labels
+Before starting any experiments with Prismer, we need to first pre-generate the modality expert labels, so we may construct a multi-label dataset. In `experts` folder, we have included all 6 experts we introduced in our paper. We have organised each expert's codebase with a shared and simple API.
+
+*Note: Specifically for segmentation experts, please first install deformable convolution operations by `cd experts/segmentation/mask2former/modeling/pixel_decoder/ops` and run `sh make.sh`.*
+
+To download pre-trained modality experts, run
+```bash
+python download_checkpoints.py --download_experts=True
+```
+
+To generate the expert labels, simply edit the `configs/experts.yaml` with the corresponding data paths, and run
+```bash
+export PYTHONPATH=.
+accelerate experts/generate_{EXPERT_NAME}.py
+```
+*Note: Expert label generation is only required for Prismer models, not for PrismerZ models.*
+
+## Experiments
+We have provided both Prismer and PrismerZ for pre-trained checkpoints (for zero-shot image captioning), as well as fined-tuned checkpoints on VQAv2 and COCO datasets. With these checkpoints, it should be expected to reproduce the exact performance listed below.
+
+| Model          | Pre-trained [Zero-shot] | COCO [Fine-tuned]   | VQAv2 [Fine-tuned] |
+|----------------|-------------------------|---------------------|-------------------|
+| PrismerZ-BASE  | COCO CIDEr [109.6]      | COCO CIDEr [133.7]	 | test-dev [76.58]  |
+| Prismer-BASE   | COCO CIDEr [122.6]      | COCO CIDEr [135.1]	 | test-dev [76.84]  |
+| PrismerZ-LARGE | COCO CIDEr [124.8]      | COCO CIDEr [135.7]	 | test-dev [77.49]  |
+| Prismer-LARGE  | COCO CIDEr [129.7]      | COCO CIDEr [136.5]	 | test-dev [78.42]  |
+
+To download pre-trained/fined-tuned checkpoints, run
+```bash
+# to download all model checkpoints (12 models in total)
+python download_checkpoints.py --download_models=True
+
+# to download specific checkpoints (Prismer-Base for fine-tuned VQA) in this example
+python download_checkpoints.py --download_models="vqa_prismer_base"
+```
+
+
+*Note: Remember to install java via `sudo apt-get install default-jre` which is required to run the official COCO caption evaluation scripts.*
+
+
+### Evaluation
+To evaluate the model checkpoints, please run
+```bash
+# zero-shot image captioning (remember to remove caption prefix in the config files)
+python train_caption.py --exp_name {MODEL_NAME} --evaluate
+
+# fine-tuned image captioning
+python train_caption.py --exp_name {MODEL_NAME} --from_checkpoint --evaluate
+
+# fine-tuned VQA
+python train_vqa.py --exp_name {MODEL_NAME} --from_checkpoint --evaluate
+```
+
+### Training / Fine-tuning
+To pre-train or fine-tune any model with or without checkpoints, please run
+```bash
+# to train/fine-tuning from scratch
+python train_{TASK}.py --exp_name {MODEL_NAME}
+
+# to train/fine-tuning from the latest checkpoints (saved every epoch)
+python train_{TASK}.py --exp_name {MODEL_NAME} --from_checkpoint 
+```
+
+We have also included model sharding in the current training script via PyTorch's official [FSDP plugin](https://pytorch.org/tutorials/intermediate/FSDP_tutorial.html). With the same training commands, additionally add `--shard_grad_op` for ZeRO-2 Sharding (Gradients + Optimiser States), or `--full_shard` for ZeRO-3 Sharding (ZeRO-2 + Network Parameters). 
+
+*Note: You should expect the error range for VQAv2 Acc. to be less than 0.1; for COCO/NoCAPs CIDEr score to be less than 1.0.*
+
+## Demo
+Finally, we have offered a minimalist example to perform image captioning in a single GPU with our fine-tuned Prismer/PrismerZ checkpoint. Simply put your images under `helpers/images` (`.jpg` images), and run
+```bash
+python demo.py --exp_name {MODEL_NAME}
+```
+
+You then can see all generated modality expert labels in the `helpers/labels` folder and the generated captions in the `helpers/images` folder.
+
+Particularly for the Prismer models, we have also offered a simple script to prettify the generated expert labels. To prettify and visualise the expert labels as well as its predicted captions, run
+```bash
+python demo_vis.py
+```
+
+*Note: Remember to set up the corresponding config in the `configs/caption.yaml` demo section. The default demo model config is for Prismer-Base.*
+
+## Citation
+
+If you found this code/work to be useful in your own research, please considering citing the following:
+
+
+```bibtex
+@article{liu2023prismer,
+    title={Prismer: A Vision-Language Model with An Ensemble of Experts},
+    author={Liu, Shikun and Fan, Linxi and Johns, Edward and Yu, Zhiding and Xiao, Chaowei and Anandkumar, Anima},
+    journal={arXiv preprint arXiv:2303.02506},
+    year={2023}
+}
+```
+
+## License
+Copyright © 2023, NVIDIA Corporation. All rights reserved.
+
+This work is made available under the Nvidia Source Code License-NC. 
+
+The model checkpoints are shared under CC-BY-NC-SA-4.0. If you remix, transform, or build upon the material, you must distribute your contributions under the same license as the original.
+
+For business inquiries, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/).
+
+## Acknowledgement
+We would like to thank all the researchers who open source their works to make this project possible. [@bjoernpl](https://github.com/bjoernpl) for contributing an automated checkpoint download script.
+
+## Contact
+If you have any questions, please contact `sk.lorenmt@gmail.com`.

prismer/dataset/__init__.py

@@ -6,12 +6,18 @@
 
 from torch.utils.data import DataLoader
 
+from dataset.pretrain_dataset import Pretrain
 from dataset.vqa_dataset import VQA
 from dataset.caption_dataset import Caption
+from dataset.classification_dataset import Classification
 
 
 def create_dataset(dataset, config):
-    if dataset == 'vqa':
+    if dataset == 'pretrain':
+        dataset = Pretrain(config)
+        return dataset
+
+    elif dataset == 'vqa':
         train_dataset = VQA(config, train=True)
         test_dataset = VQA(config, train=False)
         return train_dataset, test_dataset
@@ -20,6 +26,11 @@ def create_dataset(dataset, config):
         train_dataset = Caption(config, train=True)
         test_dataset = Caption(config, train=False)
         return train_dataset, test_dataset
+    
+    elif dataset == 'classification':
+        train_dataset = Classification(config, train=True)
+        test_dataset = Classification(config, train=False)
+        return train_dataset, test_dataset
 
 
 def create_loader(dataset, batch_size, num_workers, train, collate_fn=None):

prismer/dataset/caption_dataset.py

@@ -50,7 +50,7 @@ class Caption(Dataset):
         elif self.dataset == 'demo':
             img_path_split = self.data_list[index]['image'].split('/')
             img_name = img_path_split[-2] + '/' + img_path_split[-1]
-            image, labels, labels_info = get_expert_labels('prismer', self.label_path, img_name, 'helpers', self.experts)
+            image, labels, labels_info = get_expert_labels('', self.label_path, img_name, 'helpers', self.experts)
 
         experts = self.transform(image, labels)
         experts = post_label_process(experts, labels_info)

prismer/dataset/classification_dataset.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+from torch.utils.data import Dataset
+from dataset.utils import *
+
+
+class Classification(Dataset):
+    def __init__(self, config, train):
+        self.data_path = config['data_path']
+        self.label_path = config['label_path']
+        self.experts = config['experts']
+        self.dataset = config['dataset']
+        self.shots = config['shots']
+        self.prefix = config['prefix']
+
+        self.train = train
+        self.transform = Transform(resize_resolution=config['image_resolution'], scale_size=[0.5, 1.0], train=True)
+
+        if train:
+            data_folders = glob.glob(f'{self.data_path}/imagenet_train/*/')
+            self.data_list = [{'image': data} for f in data_folders for data in glob.glob(f + '*.JPEG')[:self.shots]]
+            self.answer_list = json.load(open(f'{self.data_path}/imagenet/' + 'imagenet_answer.json'))
+            self.class_list = json.load(open(f'{self.data_path}/imagenet/' + 'imagenet_class.json'))
+        else:
+            data_folders = glob.glob(f'{self.data_path}/imagenet/*/')
+            self.data_list = [{'image': data} for f in data_folders for data in glob.glob(f + '*.JPEG')]
+            self.answer_list = json.load(open(f'{self.data_path}/imagenet/' + 'imagenet_answer.json'))
+            self.class_list = json.load(open(f'{self.data_path}/imagenet/' + 'imagenet_class.json'))
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        img_path = self.data_list[index]['image']
+        if self.train:
+            img_path_split = img_path.split('/')
+            img_name = img_path_split[-2] + '/' + img_path_split[-1]
+            class_name = img_path_split[-2]
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, img_name, 'imagenet_train', self.experts)
+        else:
+            img_path_split = img_path.split('/')
+            img_name = img_path_split[-2] + '/' + img_path_split[-1]
+            class_name = img_path_split[-2]
+            image, labels, labels_info = get_expert_labels(self.data_path, self.label_path, img_name, 'imagenet', self.experts)
+
+        experts = self.transform(image, labels)
+        experts = post_label_process(experts, labels_info)
+
+        if self.train:
+            caption = self.prefix + ' ' + self.answer_list[int(self.class_list[class_name])].lower()
+            return experts, caption
+        else:
+            return experts, self.class_list[class_name]
+
+
+
+
+
+# import os
+# import glob
+#
+# data_path = '/Users/shikunliu/Documents/dataset/mscoco/mscoco'
+#
+# data_folders = glob.glob(f'{data_path}/*/')
+# data_list = [data for f in data_folders for data in glob.glob(f + '*.jpg')]
+
+

prismer/dataset/pretrain_dataset.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import glob
+
+from torch.utils.data import Dataset
+from dataset.utils import *
+
+
+class Pretrain(Dataset):
+    def __init__(self, config):
+        self.cc12m_data_path = config['cc12m_data_path']
+        self.cc3m_data_path = config['cc3m_data_path']
+        self.coco_data_path = config['coco_data_path']
+        self.vg_data_path = config['vg_data_path']
+        self.label_path = config['label_path']
+        self.experts = config['experts']
+
+        self.data_list = []
+        if 'cc12m' in config['datasets']:
+            data_folders = glob.glob(f'{self.cc12m_data_path}/cc12m/*/')
+            self.data_list += [{'image': data} for f in data_folders for data in glob.glob(f + '*.jpg')]
+        if 'cc3m_sgu' in config['datasets']:
+            data_folders = glob.glob(f'{self.cc3m_data_path}/cc3m_sgu/*/')
+            self.data_list += [{'image': data} for f in data_folders for data in glob.glob(f + '*.jpg')]
+        if 'coco' in config['datasets']:
+            self.data_list += json.load(open(os.path.join(self.coco_data_path, 'coco_karpathy_train.json'), 'r'))
+        if 'vg' in config['datasets']:
+            self.data_list += json.load(open(os.path.join(self.vg_data_path, 'vg_caption.json'), 'r'))
+
+        self.transform = Transform(resize_resolution=config['image_resolution'], scale_size=[0.5, 1.5], train=True)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        img_path = self.data_list[index]['image']
+
+        if 'cc12m' in img_path:
+            img_path_split = img_path.split('/')
+            img_name = img_path_split[-2] + '/' + img_path_split[-1]
+            image, labels, labels_info = get_expert_labels(self.cc12m_data_path, self.label_path, img_name, 'cc12m', self.experts)
+
+            caption_path = img_path.replace('.jpg', '.txt')
+            with open(caption_path) as f:
+                caption = f.readlines()[0]
+                
+        elif 'cc3m_sgu' in img_path:
+            img_path_split = img_path.split('/')
+            img_name = img_path_split[-2] + '/' + img_path_split[-1]
+            image, labels, labels_info = get_expert_labels(self.cc3m_data_path, self.label_path, img_name, 'cc3m_sgu', self.experts)
+
+            caption_path = img_path.replace('.jpg', '.txt')
+            with open(caption_path) as f:
+                caption = f.readlines()[0]
+
+        elif 'train2014' in img_path or 'val2014' in img_path:
+            image, labels, labels_info = get_expert_labels(self.coco_data_path, self.label_path, img_path, 'vqav2', self.experts)
+            caption = self.data_list[index]['caption']
+
+        elif 'visual-genome' in img_path:
+            img_path_split = img_path.split('/')
+            img_name = img_path_split[-2] + '/' + img_path_split[-1]
+            image, labels, labels_info = get_expert_labels(self.vg_data_path, self.label_path, img_name, 'vg', self.experts)
+            caption = self.data_list[index]['caption']
+
+        experts = self.transform(image, labels)
+        experts = post_label_process(experts, labels_info)
+        caption = pre_caption(caption, max_words=30)
+        return experts, caption

prismer/dataset/utils.py

@@ -12,16 +12,12 @@ import PIL.Image as Image
 import numpy as np
 import torchvision.transforms as transforms
 import torchvision.transforms.functional as transforms_f
-import pathlib
 from dataset.randaugment import RandAugment
 
-
-cur_dir = pathlib.Path(__file__).parent
-
-COCO_FEATURES = torch.load(cur_dir / 'coco_features.pt')['features']
-ADE_FEATURES = torch.load(cur_dir / 'ade_features.pt')['features']
-DETECTION_FEATURES = torch.load(cur_dir / 'detection_features.pt')['features']
-BACKGROUND_FEATURES = torch.load(cur_dir / 'background_features.pt')
+COCO_FEATURES = torch.load('dataset/coco_features.pt')['features']
+ADE_FEATURES = torch.load('dataset/ade_features.pt')['features']
+DETECTION_FEATURES = torch.load('dataset/detection_features.pt')['features']
+BACKGROUND_FEATURES = torch.load('dataset/background_features.pt')
 
 
 class Transform:

prismer/dataset/vqa_dataset.py

@@ -6,8 +6,6 @@
 
 from torch.utils.data import Dataset
 from dataset.utils import *
-from PIL import ImageFile
-ImageFile.LOAD_TRUNCATED_IMAGES = True
 
 
 class VQA(Dataset):

prismer/experts/generate_depth.py

@@ -20,7 +20,7 @@ from tqdm import tqdm
 model, transform = load_expert_model(task='depth')
 accelerator = Accelerator(mixed_precision='fp16')
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'depth')
 

prismer/experts/generate_edge.py

@@ -22,7 +22,7 @@ from tqdm import tqdm
 model, transform = load_expert_model(task='edge')
 accelerator = Accelerator(mixed_precision='fp16')
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'edge')
 

prismer/experts/generate_normal.py

@@ -22,7 +22,7 @@ import numpy as np
 model, transform = load_expert_model(task='normal')
 accelerator = Accelerator(mixed_precision='fp16')
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'normal')
 

prismer/experts/generate_objdet.py

@@ -22,7 +22,7 @@ from tqdm import tqdm
 model, transform = load_expert_model(task='obj_detection')
 accelerator = Accelerator(mixed_precision='fp16')
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = config['save_path']
 

prismer/experts/generate_ocrdet.py

@@ -26,7 +26,7 @@ model, transform = load_expert_model(task='ocr_detection')
 accelerator = Accelerator(mixed_precision='fp16')
 pca_clip = pk.load(open('dataset/clip_pca.pkl', 'rb'))
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'ocr_detection')
 

prismer/experts/generate_segmentation.py

@@ -20,7 +20,7 @@ from tqdm import tqdm
 model, transform = load_expert_model(task='seg_coco')
 accelerator = Accelerator(mixed_precision='fp16')
 
-config = yaml.load(open('prismer/configs/experts.yaml', 'r'), Loader=yaml.Loader)
+config = yaml.load(open('configs/experts.yaml', 'r'), Loader=yaml.Loader)
 data_path = config['data_path']
 save_path = os.path.join(config['save_path'], 'seg_coco')
 

prismer/model/prismer.py

@@ -5,15 +5,12 @@
 # https://github.com/NVlabs/prismer/blob/main/LICENSE
 
 import json
-import pathlib
 import torch.nn as nn
 
 from model.modules.vit import load_encoder
 from model.modules.roberta import load_decoder
 from transformers import RobertaTokenizer, RobertaConfig
 
-cur_dir = pathlib.Path(__file__).parent
-
 
 class Prismer(nn.Module):
     def __init__(self, config):
@@ -29,7 +26,7 @@ class Prismer(nn.Module):
             elif exp in ['obj_detection', 'ocr_detection']:
                 self.experts[exp] = 64
 
-        prismer_config = json.load(open(f'{cur_dir.parent}/configs/prismer.json', 'r'))[config['prismer_model']]
+        prismer_config = json.load(open('configs/prismer.json', 'r'))[config['prismer_model']]
         roberta_config = RobertaConfig.from_dict(prismer_config['roberta_model'])
 
         self.tokenizer = RobertaTokenizer.from_pretrained(prismer_config['roberta_model']['model_name'])

prismer/requirements.txt

@@ -0,0 +1,19 @@
+git+https://github.com/openai/CLIP.git
+git+https://github.com/facebookresearch/detectron2.git@5aeb252b194b93dc2879b4ac34bc51a31b5aee13
+accelerate
+fairscale
+timm
+transformers
+einops
+scikit-learn==0.24.2
+pycocoevalcap
+editdistance
+shapely
+pyclipper
+yacs
+pycocotools
+geffnet
+fire
+huggingface_hub
+rich
+ruamel.yaml

prismer/train_caption.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import argparse
+import numpy as np
+import random
+import time
+import functools
+import json
+import torch
+import os
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
+from model.prismer_caption import PrismerCaption
+from model.modules.utils import interpolate_pos_embed
+from dataset import create_dataset, create_loader
+from utils import *
+from tqdm import tqdm
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mode', default='')
+parser.add_argument('--port', default='')
+
+parser.add_argument('--config', default='configs/caption.yaml')
+parser.add_argument('--from_checkpoint', action='store_true')
+parser.add_argument('--evaluate', action='store_true')
+parser.add_argument('--target_dataset', default='coco', type=str)
+parser.add_argument('--shard_grad_op', action='store_true')
+parser.add_argument('--full_shard', action='store_true')
+parser.add_argument('--exp_name', default='', type=str)
+parser.add_argument('--mixed_precision', default='fp16', type=str)
+parser.add_argument('--seed', default=42, type=int)
+args = parser.parse_args()
+
+config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)[args.target_dataset]
+torch.manual_seed(args.seed)
+np.random.seed(args.seed)
+random.seed(args.seed)
+
+train_dataset, test_dataset = create_dataset('caption', config)
+train_loader = create_loader(train_dataset, batch_size=config['batch_size_train'], num_workers=8, train=True)
+test_loader = create_loader(test_dataset, batch_size=config['batch_size_test'], num_workers=8, train=False)
+
+
+model = PrismerCaption(config)
+tokenizer = model.tokenizer
+
+if args.shard_grad_op:  # Model Sharding: ZeRO 2
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+
+elif args.full_shard:  # Model Sharding: ZeRO 3
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    from model.modules.vit import ResidualAttentionBlock
+    from model.modules.resampler import PerceiverAttentionBlock
+    from model.modules.roberta import RobertaLayer
+    auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={
+            ResidualAttentionBlock,
+            PerceiverAttentionBlock,
+            RobertaLayer
+        },
+    )
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.FULL_SHARD,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 auto_wrap_policy=auto_wrap_policy,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+else:
+    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+
+# Reload saved states
+if not args.from_checkpoint:
+    state_dict = torch.load(f'logging/pretrain_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    state_dict['expert_encoder.positional_embedding'] = interpolate_pos_embed(state_dict['expert_encoder.positional_embedding'],
+                                                                              len(model.expert_encoder.positional_embedding))
+    model.load_state_dict(state_dict)
+    start_epoch = 0
+else:
+    state_dict = torch.load(f'logging/caption_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    if os.path.exists(f'logging/caption_{args.exp_name}/epoch.pt'):
+        start_epoch = torch.load(f'logging/caption_{args.exp_name}/epoch.pt')[0] + 1
+    else:
+        start_epoch = 0
+    model.load_state_dict(state_dict)
+    accelerator.print(f'Start re-training from checkpoint with Epoch {start_epoch}')
+
+optimizer = torch.optim.AdamW(params=filter(lambda p: p.requires_grad, model.parameters()),
+                              lr=config['init_lr'], weight_decay=config['weight_decay'])
+
+if args.shard_grad_op or args.full_shard:
+    optimizer, train_loader, test_loader = accelerator.prepare(optimizer, train_loader, test_loader)
+else:
+    model, optimizer, train_loader, test_loader = accelerator.prepare(model, optimizer, train_loader, test_loader)
+
+best = 0
+start_time = time.time()
+if not args.evaluate:
+    for epoch in range(start_epoch, config['max_epoch']):
+        train_loss = 0
+        num_train_elems = 0
+        model.train()
+        for i, (experts, caption) in enumerate(tqdm(train_loader)):
+            cosine_lr_schedule(optimizer, epoch * len(train_loader) + i, config['max_epoch'] * len(train_loader), config['init_lr'], config['min_lr'])
+            
+            loss = model(experts, caption, prefix=config['prefix'])
+
+            optimizer.zero_grad()
+            accelerator.backward(loss)
+            optimizer.step()
+
+            train_loss += loss.item()
+            num_train_elems += 1
+
+        model.eval()
+        result = []
+        with torch.no_grad():
+            for step, (experts, data_ids) in enumerate(tqdm(test_loader)):
+                captions = model(experts, train=False, prefix=config['prefix'])
+
+                if accelerator.use_distributed:
+                    captions = tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
+                    captions = captions.to(experts['rgb'].device)
+                    data_ids, captions = accelerator.gather_for_metrics((data_ids, captions))
+
+                    for data_id, caption in zip(data_ids, captions):
+                        caption = tokenizer.decode(caption, skip_special_tokens=True)
+                        if args.target_dataset == 'coco':
+                            image_id = int(test_loader.dataset.data_list[data_id]['image'].split('/')[-1].strip('.jpg').split('_')[-1])
+                            result.append({"image_id": image_id, "caption": caption.capitalize() + '.'})
+                        elif args.target_dataset == 'nocaps':
+                            result.append({"image_id": test_loader.dataset.data_list[data_id]['img_id'],
+                                           "caption": caption.capitalize() + '.'})
+
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            json.dump(result, open(f'/results/caption_results_{args.exp_name}_{args.target_dataset}.json', 'w'))
+            if args.target_dataset == 'coco':
+                coco_eval = coco_caption_eval(f'{config["data_path"]}/coco_karpathy_test_gt.json', result)
+                torch.save([coco_eval.eval['CIDEr']], f'logging/caption_{args.exp_name}/temp_cider.pt')
+                if not os.path.isfile(f'logging/caption_{args.exp_name}/cider.pt'):
+                    torch.save([coco_eval.eval['CIDEr']], f'logging/caption_{args.exp_name}/cider.pt')
+
+        accelerator.wait_for_everyone()
+        cider = torch.load(f'logging/caption_{args.exp_name}/cider.pt')[0]
+        curr_cider = torch.load(f'logging/caption_{args.exp_name}/temp_cider.pt')[0]
+
+        if cider < curr_cider:
+            train_loss /= num_train_elems
+            accelerator.print(f"Epoch {epoch:03d} | loss: {train_loss:.4f} || Time: {(time.time() - start_time):.4f}")
+            accelerator.save_state(f'logging/caption_{args.exp_name}')
+            accelerator.save([epoch], f'logging/caption_{args.exp_name}/epoch.pt')
+            accelerator.save([curr_cider], f'logging/caption_{args.exp_name}/cider.pt')
+
+
+model.eval()
+if accelerator.is_main_process:
+    result = []
+
+with torch.no_grad():
+    for step, (experts, data_ids) in enumerate(tqdm(test_loader)):
+        captions = model(experts, train=False, prefix=config['prefix'])
+
+        if accelerator.use_distributed:
+            captions = tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
+            captions = captions.to(experts['rgb'].device)
+            data_ids, captions = accelerator.gather_for_metrics((data_ids, captions))
+
+        if accelerator.is_main_process:
+            for data_id, caption in zip(data_ids, captions):
+                caption = tokenizer.decode(caption, skip_special_tokens=True)
+                if args.target_dataset == 'coco':
+                    image_id = int(test_loader.dataset.data_list[data_id]['image'].split('/')[-1].strip('.jpg').split('_')[-1])
+                    result.append({"image_id": image_id, "caption": caption.capitalize() + '.'})
+                elif args.target_dataset == 'nocaps':
+                    result.append({"image_id": test_loader.dataset.data_list[data_id]['img_id'],
+                                   "caption": caption.capitalize() + '.'})
+
+accelerator.wait_for_everyone()
+if accelerator.is_main_process:
+    json.dump(result, open(f'/results/caption_results_{args.exp_name}_{args.target_dataset}.json', 'w'))
+    if args.target_dataset == 'coco':
+        coco_caption_eval(f'{config["data_path"]}/coco_karpathy_test_gt.json', result)
+
+

prismer/train_classification.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import argparse
+import numpy as np
+import random
+import time
+import functools
+import torch
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
+from model.prismer_caption import PrismerCaption
+from model.modules.utils import interpolate_pos_embed
+from dataset import create_dataset, create_loader
+from tqdm import tqdm
+from utils import *
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mode', default='')
+parser.add_argument('--port', default='')
+
+parser.add_argument('--config', default='configs/classification.yaml')
+parser.add_argument('--from_checkpoint', action='store_true')
+parser.add_argument('--evaluate', action='store_true')
+parser.add_argument('--exp_name', default='', type=str)
+parser.add_argument('--shard_grad_op', action='store_true')
+parser.add_argument('--full_shard', action='store_true')
+parser.add_argument('--mixed_precision', default='fp16', type=str)
+parser.add_argument('--seed', default=42, type=int)
+args = parser.parse_args()
+
+config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)
+torch.manual_seed(args.seed)
+np.random.seed(args.seed)
+random.seed(args.seed)
+
+train_dataset, test_dataset = create_dataset('classification', config)
+train_loader = create_loader(train_dataset, batch_size=config['batch_size_train'], num_workers=8, train=True)
+test_loader = create_loader(test_dataset, batch_size=config['batch_size_test'], num_workers=8, train=False)
+model = PrismerCaption(config)
+
+if args.shard_grad_op:  # Model Sharding: ZeRO 2
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+
+elif args.full_shard:  # Model Sharding: ZeRO 3
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    from model.modules.vit import ResidualAttentionBlock
+    from model.modules.resampler import PerceiverAttentionBlock
+    from model.modules.roberta import RobertaLayer
+    auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={
+            ResidualAttentionBlock,
+            PerceiverAttentionBlock,
+            RobertaLayer
+        },
+    )
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.FULL_SHARD,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 auto_wrap_policy=auto_wrap_policy,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+else:
+    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+
+# Reload saved states
+if not args.from_checkpoint:
+    state_dict = torch.load(f'logging/pretrain_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    state_dict['expert_encoder.positional_embedding'] = interpolate_pos_embed(state_dict['expert_encoder.positional_embedding'],
+                                                                              len(model.expert_encoder.positional_embedding))
+    model.load_state_dict(state_dict)
+    start_epoch = 0
+else:
+    state_dict = torch.load(f'logging/classification_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    if os.path.exists(f'logging/classification_{args.exp_name}/epoch.pt'):
+        start_epoch = torch.load(f'logging/classification_{args.exp_name}/epoch.pt')[0] + 1
+    else:
+        start_epoch = 0
+    model.load_state_dict(state_dict)
+    accelerator.print(f'Start re-training from checkpoint with Epoch {start_epoch}')
+
+optimizer = torch.optim.AdamW(params=filter(lambda p: p.requires_grad, model.parameters()),
+                              lr=config['init_lr'], weight_decay=config['weight_decay'])
+
+if args.shard_grad_op or args.full_shard:
+    optimizer, train_loader, test_loader = accelerator.prepare(optimizer, train_loader, test_loader)
+else:
+    model, optimizer, train_loader, test_loader = accelerator.prepare(model, optimizer, train_loader, test_loader)
+
+start_time = time.time()
+best = 0
+for epoch in range(start_epoch, config['max_epoch']):
+    train_loss = 0
+    num_train_elems = 0
+    model.train()
+    for i, (experts, caption) in enumerate(tqdm(train_loader)):
+        cosine_lr_schedule(optimizer, epoch * len(train_loader) + i, config['max_epoch'] * len(train_loader), config['init_lr'], config['min_lr'])
+        loss = model(experts, caption, prefix=config['prefix'])
+
+        optimizer.zero_grad()
+        accelerator.backward(loss)
+        optimizer.step()
+
+        train_loss += loss.item()
+        num_train_elems += 1
+        
+    train_loss /= num_train_elems
+    accelerator.print(f"Epoch {epoch:03d} | loss: {train_loss:.4f} || Time: {(time.time() - start_time):.4f}")
+    
+    if (epoch + 1) % 5 == 0:
+        model.eval()
+        num_test_elems = 0
+        accurate = 0
+        with torch.no_grad():
+            answer_list = test_loader.dataset.answer_list
+            for step, (experts, gt) in enumerate(tqdm(test_loader)):
+                predictions = model(experts, answer=answer_list, train=False, prefix=config['prefix'], k_test=config['k_test'], inference='rank')
+    
+                if accelerator.use_distributed:
+                    predictions, gt = accelerator.gather_for_metrics((predictions, gt))
+
+                accurate_preds = predictions == gt
+                num_test_elems += accurate_preds.shape[0]
+                accurate += accurate_preds.long().sum()
+            eval_metric = accurate.item() / num_test_elems
+       
+        accelerator.wait_for_everyone()
+        accelerator.print(f'{config["shots"]}-Shot Acc: {eval_metric}')
+    
+        if eval_metric > best:
+            best = eval_metric
+            accelerator.save_state(f'logging/classification_{args.exp_name}')
+            accelerator.save([epoch], f'logging/classification_{args.exp_name}/epoch.pt')
+
+
+
+
+
+
+
+
+

prismer/train_pretrain.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import argparse
+import numpy as np
+import random
+import time
+import datetime
+import functools
+import torch
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
+from model.prismer_caption import PrismerCaption
+from dataset import create_dataset, create_loader
+from utils import *
+from tqdm import tqdm
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mode', default='')
+parser.add_argument('--port', default='')
+
+parser.add_argument('--config', default='configs/pretrain.yaml')
+parser.add_argument('--from_checkpoint', action='store_true')
+parser.add_argument('--shard_grad_op', action='store_true')
+parser.add_argument('--full_shard', action='store_true')
+parser.add_argument('--exp_name', default='', type=str)
+parser.add_argument('--mixed_precision', default='fp16', type=str)
+parser.add_argument('--seed', default=42, type=int)
+args = parser.parse_args()
+
+config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)
+torch.manual_seed(args.seed)
+np.random.seed(args.seed)
+random.seed(args.seed)
+
+train_dataset = create_dataset('pretrain', config)
+train_loader = create_loader(train_dataset, batch_size=config['batch_size_train'], num_workers=8, train=True)
+
+model = PrismerCaption(config)
+if args.shard_grad_op:  # Model Sharding: ZeRO 2
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType, CPUOffload
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 cpu_offload=CPUOffload(offload_params=False),
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+
+elif args.full_shard:  # Model Sharding: ZeRO 3
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    from model.modules.vit import ResidualAttentionBlock
+    from model.modules.resampler import PerceiverAttentionBlock
+    from model.modules.roberta import RobertaLayer
+    auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={
+            ResidualAttentionBlock,
+            PerceiverAttentionBlock,
+            RobertaLayer
+        },
+    )
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.FULL_SHARD,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 auto_wrap_policy=auto_wrap_policy,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+else:
+    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+
+# Reload saved states
+if args.from_checkpoint:
+    state_dict = torch.load(f'logging/pretrain_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    if os.path.exists(f'logging/pretrain_{args.exp_name}/epoch.pt'):
+        start_epoch = torch.load(f'logging/pretrain_{args.exp_name}/epoch.pt')[0] + 1
+    else:
+        start_epoch = 0
+    model.load_state_dict(state_dict)
+    accelerator.print(f'Start re-training from checkpoint with Epoch {start_epoch}')
+else:
+    start_epoch = 0
+
+optimizer = torch.optim.AdamW(params=filter(lambda p: p.requires_grad, model.parameters()),
+                              lr=config['init_lr'], weight_decay=config['weight_decay'])
+
+if args.shard_grad_op or args.full_shard:
+    optimizer, train_loader = accelerator.prepare(optimizer, train_loader)
+else:
+    model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader)
+
+
+start_time = time.time()
+warmup_step = 0
+for epoch in range(start_epoch, config['max_epoch']):
+    cosine_lr_schedule(optimizer, epoch, config['max_epoch'], config['init_lr'], config['min_lr'])
+
+    train_loss = 0
+    num_train_elems = 0
+    model.train()
+    for i, (experts, caption) in enumerate(tqdm(train_loader)):
+        if warmup_step < config['warmup_steps']:
+            warmup_lr_schedule(optimizer, warmup_step, config['warmup_steps'], config['warmup_lr'], config['init_lr'])
+            warmup_step += 1
+
+        loss = model(experts, caption)
+
+        optimizer.zero_grad()
+        accelerator.backward(loss)
+        optimizer.step()
+
+        train_loss += loss.item()
+        num_train_elems += 1
+
+    train_loss /= num_train_elems
+    accelerator.print(f"Epoch {epoch:03d} | loss: {train_loss:.4f} || Time: {(time.time() - start_time):.4f}")
+    accelerator.save_state(f'logging/pretrain_{args.exp_name}')
+    accelerator.save([epoch], f'logging/pretrain_{args.exp_name}/epoch.pt')
+
+total_time = time.time() - start_time
+total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+accelerator.print('Training time {}'.format(total_time_str))
+
+

prismer/train_vqa.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2023, NVIDIA Corporation & Affiliates. All rights reserved.
+#
+# This work is made available under the Nvidia Source Code License-NC.
+# To view a copy of this license, visit
+# https://github.com/NVlabs/prismer/blob/main/LICENSE
+
+import argparse
+import numpy as np
+import random
+import time
+import datetime
+import functools
+import torch
+try:
+    import ruamel_yaml as yaml
+except ModuleNotFoundError:
+    import ruamel.yaml as yaml
+    
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
+from model.prismer_vqa import PrismerVQA
+from model.modules.utils import interpolate_pos_embed
+from dataset import create_dataset, create_loader
+from utils import *
+from tqdm import tqdm
+import json
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--mode', default='')
+parser.add_argument('--port', default='')
+
+parser.add_argument('--config', default='configs/vqa.yaml')
+parser.add_argument('--from_checkpoint', action='store_true')
+parser.add_argument('--evaluate', action='store_true')
+parser.add_argument('--exp_name', default='', type=str)
+parser.add_argument('--shard_grad_op', action='store_true')
+parser.add_argument('--full_shard', action='store_true')
+parser.add_argument('--mixed_precision', default='fp16', type=str)
+parser.add_argument('--seed', default=42, type=int)
+args = parser.parse_args()
+
+config = yaml.load(open(args.config, 'r'), Loader=yaml.Loader)
+torch.manual_seed(args.seed)
+np.random.seed(args.seed)
+random.seed(args.seed)
+
+train_dataset, test_dataset = create_dataset('vqa', config)
+train_loader = create_loader(train_dataset, batch_size=config['batch_size_train'], num_workers=8, train=True)
+test_loader = create_loader(test_dataset, batch_size=config['batch_size_test'], num_workers=8, train=False)
+
+model = PrismerVQA(config)
+tokenizer = model.tokenizer
+
+if args.shard_grad_op:  # Model Sharding: ZeRO 2
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+
+elif args.full_shard:  # Model Sharding: ZeRO 3
+    from torch.distributed.fsdp import MixedPrecision, BackwardPrefetch, ShardingStrategy, StateDictType
+    from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+    from model.modules.vit import ResidualAttentionBlock
+    from model.modules.resampler import PerceiverAttentionBlock
+    from model.modules.roberta import RobertaLayer
+    auto_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls={
+            ResidualAttentionBlock,
+            PerceiverAttentionBlock,
+            RobertaLayer
+        },
+    )
+    fsdp_plugin = FullyShardedDataParallelPlugin(sharding_strategy=ShardingStrategy.FULL_SHARD,
+                                                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+                                                 mixed_precision_policy=MixedPrecision(param_dtype=torch.float16,
+                                                                                       reduce_dtype=torch.float16,
+                                                                                       buffer_dtype=torch.float16),
+                                                 state_dict_type=StateDictType.FULL_STATE_DICT,
+                                                 auto_wrap_policy=auto_wrap_policy,
+                                                 ignored_modules=model.ignored_modules)
+    accelerator = Accelerator(mixed_precision=args.mixed_precision, fsdp_plugin=fsdp_plugin)
+    model = accelerator.prepare(model)
+else:
+    accelerator = Accelerator(mixed_precision=args.mixed_precision)
+
+# Reload saved states
+if not args.from_checkpoint:
+    state_dict = torch.load(f'logging/pretrain_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    state_dict['expert_encoder.positional_embedding'] = interpolate_pos_embed(state_dict['expert_encoder.positional_embedding'],
+                                                                              len(model.expert_encoder.positional_embedding))
+    model.load_state_dict(state_dict)
+    start_epoch = 0
+else:
+    state_dict = torch.load(f'logging/vqa_{args.exp_name}/pytorch_model.bin', map_location='cpu')
+    if os.path.exists(f'logging/vqa_{args.exp_name}/epoch.pt'):
+        start_epoch = torch.load(f'logging/vqa_{args.exp_name}/epoch.pt')[0] + 1
+    else:
+        start_epoch = 0
+    model.load_state_dict(state_dict)
+    accelerator.print(f'Start re-training from checkpoint with Epoch {start_epoch}')
+
+optimizer = torch.optim.AdamW(params=filter(lambda p: p.requires_grad, model.parameters()),
+                              lr=config['init_lr'], weight_decay=config['weight_decay'])
+
+if args.shard_grad_op or args.full_shard:
+    optimizer, train_loader, test_loader = accelerator.prepare(optimizer, train_loader, test_loader)
+else:
+    model, optimizer, train_loader, test_loader = accelerator.prepare(model, optimizer, train_loader, test_loader)
+
+start_time = time.time()
+if not args.evaluate:
+    for epoch in range(start_epoch, config['max_epoch']):
+        cosine_lr_schedule(optimizer, epoch, config['max_epoch'], config['init_lr'], config['min_lr'])
+
+        train_loss = 0
+        num_train_elems = 0
+        model.train()
+        for i, (experts, question, answer, weights) in enumerate(tqdm(train_loader)):
+            loss = model(experts, question, answer, train=True, weights=weights)
+            optimizer.zero_grad()
+            accelerator.backward(loss)
+            optimizer.step()
+
+            train_loss += loss.item()
+            num_train_elems += 1
+
+        train_loss /= num_train_elems
+        accelerator.print(f"Epoch {epoch:03d} | loss: {train_loss:.4f} || Time: {(time.time() - start_time):.4f}")
+        accelerator.save_state(f'logging/vqa_{args.exp_name}')
+        accelerator.save([epoch], f'logging/vqa_{args.exp_name}/epoch.pt')
+
+model.eval()
+if accelerator.is_main_process:
+    result = []
+
+with torch.no_grad():
+    if config['inference'] == 'rank':
+        answer_list = test_loader.dataset.answer_list
+
+    for step, (experts, data_ids, question, question_id) in enumerate(tqdm(test_loader)):
+        if config['inference'] == 'generate':
+            answers = model(experts, question, train=False, inference='generate')
+
+            if accelerator.use_distributed:
+                answers = tokenizer(answers, max_length=15, padding='max_length', return_tensors='pt').input_ids
+                answers = answers.to(experts['rgb'].device)
+                data_ids, answers, question_id = accelerator.gather_for_metrics((data_ids, answers, question_id))
+
+            if accelerator.is_main_process:
+                for data_id, answer, ques_id in zip(data_ids, answers, question_id):
+                    answer = tokenizer.decode(answer, skip_special_tokens=True)
+                    result.append({"question_id": int(ques_id.item()), "answer": answer})
+
+        elif config['inference'] == 'rank':
+            answer_ids = model(experts, question, answer_list, train=False, inference='rank', k_test=config['k_test'])
+
+            if accelerator.use_distributed:
+                answer_ids, question_id = accelerator.gather_for_metrics((answer_ids, question_id))
+
+            if accelerator.is_main_process:
+                for ques_id, answer_id in zip(question_id, answer_ids):
+                    result.append({"question_id": int(ques_id.item()), "answer": answer_list[answer_id]})
+
+
+accelerator.wait_for_everyone()
+if accelerator.is_main_process:
+    json.dump(result, open(f'/results/vqa_results_{args.exp_name}.json', 'w'))
+
+
+total_time = time.time() - start_time
+total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+accelerator.print('Training time {}'.format(total_time_str))
+
+

prismer_model.py

@@ -20,22 +20,32 @@ from model.prismer_caption import PrismerCaption
 
 def download_models() -> None:
     if not pathlib.Path('prismer/experts/expert_weights/').exists():
-        subprocess.run(shlex.split('python download_checkpoints.py --download_experts=True'), cwd='prismer')
-
+        subprocess.run(shlex.split(
+            'python download_checkpoints.py --download_experts=True'),
+                       cwd='prismer')
     model_names = [
-        # 'vqa_prismer_base',
-        # 'vqa_prismer_large',
+        'vqa_prismer_base',
+        'vqa_prismer_large',
+        'vqa_prismerz_base',
+        'vqa_prismerz_large',
+        'caption_prismerz_base',
+        'caption_prismerz_large',
         'caption_prismer_base',
         'caption_prismer_large',
     ]
     for model_name in model_names:
         if pathlib.Path(f'prismer/logging/{model_name}').exists():
             continue
-        subprocess.run(shlex.split(f'python download_checkpoints.py --download_models={model_name}'), cwd='prismer')
+        subprocess.run(shlex.split(
+            f'python download_checkpoints.py --download_models={model_name}'),
+                       cwd='prismer')
 
 
 def build_deformable_conv() -> None:
-    subprocess.run(shlex.split('sh make.sh'), cwd='prismer/experts/segmentation/mask2former/modeling/pixel_decoder/ops')
+    subprocess.run(
+        shlex.split('sh make.sh'),
+        cwd=
+        'prismer/experts/segmentation/mask2former/modeling/pixel_decoder/ops')
 
 
 def run_experts(image_path: str) -> tuple[str | None, ...]:
@@ -46,18 +56,40 @@ def run_experts(image_path: str) -> tuple[str | None, ...]:
     out_path = image_dir / 'image.jpg'
     cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
 
-    expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
+    expert_names = [
+        'depth',
+        'edge',
+        'normal',
+        'objdet',
+        'ocrdet',
+        'segmentation',
+    ]
     for expert_name in expert_names:
         env = os.environ.copy()
         if 'PYTHONPATH' in env:
             env['PYTHONPATH'] = f'{submodule_dir.as_posix()}:{env["PYTHONPATH"]}'
         else:
             env['PYTHONPATH'] = submodule_dir.as_posix()
-        subprocess.run(shlex.split(f'python experts/generate_{expert_name}.py'), cwd='prismer', env=env, check=True)
-
-    keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
-    results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]
-    return tuple(path.as_posix() if path.exists() else None for path in results)
+        subprocess.run(
+            shlex.split(f'python experts/generate_{expert_name}.py'),
+            cwd='prismer',
+            env=env,
+            check=True)
+
+    keys = [
+        'depth',
+        'edge',
+        'normal',
+        'seg_coco',
+        'obj_detection',
+        'ocr_detection',
+    ]
+    results = [
+        pathlib.Path('prismer/helpers/labels') / key /
+        'helpers/images/image.png' for key in keys
+    ]
+    return tuple(path.as_posix() if path.exists() else None
+                 for path in results)
 
 
 class Model:
@@ -71,14 +103,28 @@ class Model:
         if exp_name == self.exp_name:
             return
         config = {
-            'dataset': 'demo',
-            'data_path': 'prismer/helpers',
-            'label_path': 'prismer/helpers/labels',
-            'experts': ['depth', 'normal', 'seg_coco', 'edge', 'obj_detection', 'ocr_detection'],
-            'image_resolution': 480,
-            'prismer_model': 'prismer_base',
-            'freeze': 'freeze_vision',
-            'prefix': 'A picture of',
+            'dataset':
+            'demo',
+            'data_path':
+            'prismer/helpers',
+            'label_path':
+            'prismer/helpers/labels',
+            'experts': [
+                'depth',
+                'normal',
+                'seg_coco',
+                'edge',
+                'obj_detection',
+                'ocr_detection',
+            ],
+            'image_resolution':
+            480,
+            'prismer_model':
+            'prismer_base',
+            'freeze':
+            'freeze_vision',
+            'prefix':
+            'A picture of',
         }
         model = PrismerCaption(config)
         state_dict = torch.load(
@@ -96,17 +142,27 @@ class Model:
     @torch.inference_mode()
     def run_caption_model(self, exp_name: str) -> str:
         self.set_model(exp_name)
+
         _, test_dataset = create_dataset('caption', self.config)
-        test_loader = create_loader(test_dataset, batch_size=1, num_workers=4, train=False)
+        test_loader = create_loader(test_dataset,
+                                    batch_size=1,
+                                    num_workers=4,
+                                    train=False)
         experts, _ = next(iter(test_loader))
-        captions = self.model(experts, train=False, prefix=self.config['prefix'])
-        captions = self.tokenizer(captions, max_length=30, padding='max_length', return_tensors='pt').input_ids
+        captions = self.model(experts,
+                              train=False,
+                              prefix=self.config['prefix'])
+        captions = self.tokenizer(captions,
+                                  max_length=30,
+                                  padding='max_length',
+                                  return_tensors='pt').input_ids
         caption = captions.to(experts['rgb'].device)[0]
         caption = self.tokenizer.decode(caption, skip_special_tokens=True)
         caption = caption.capitalize() + '.'
         return caption
 
-    def run_caption(self, image_path: str, model_name: str) -> tuple[str | None, ...]:
+    def run_caption(self, image_path: str,
+                    model_name: str) -> tuple[str | None, ...]:
         out_paths = run_experts(image_path)
-        # caption = self.run_caption_model(model_name)
-        return None, *out_paths
+        caption = self.run_caption_model(model_name)
+        return caption, *out_paths