Add files using upload-large-folder tool

multimodal/examples/albef/configs/retrieval.yaml

@@ -0,0 +1,73 @@
+hidden_size: &hidden_size 768
+vocab_size: &vocab_size 30522
+type_vocab_size: &type_vocab_size 2
+max_position_embeddings: &max_position_embeddings 512
+pad_token_id: &pad_token_id 0
+embed_size: &embed_size 256
+
+seed: 42
+world_size: 1
+device: "cuda"
+dist_url: "env://"
+output_path: "./examples/albef/outputs/retrieval_output.pt"
+
+datamodule_args:
+  train_files: ["./examples/albef/data_files/coco_train.json"]
+  test_files: ["./examples/albef/data_files/coco_test.json"]
+  image_root: "./examples/albef/data_files/coco"
+  batch_size: 32
+  num_workers: 8
+
+vision_encoder_args:
+  hidden_size: *hidden_size
+  image_size: 384
+  patch_size: 16
+  num_hidden_layers: 12
+  num_attention_heads: 12
+  mlp_dim: 3072
+  dropout: 0.0
+  attention_dropout: 0.0
+  layer_norm_eps: 1e-6
+
+text_encoder_args:
+  vocab_size: *vocab_size
+  hidden_size: *hidden_size
+  type_vocab_size: *type_vocab_size
+  max_position_embeddings: *max_position_embeddings
+  pad_token_id: *pad_token_id
+  num_hidden_layers: 6
+  num_attention_heads: 12
+  intermediate_size: 3072
+  layer_norm_eps: 1e-12
+  dropout: 0.0
+
+multimodal_encoder_args:
+  hidden_size: *hidden_size
+  num_hidden_layers: 6
+  num_attention_heads: 12
+  intermediate_size: 3072
+  layer_norm_eps: 1e-12
+
+projection_args:
+  in_features: *hidden_size
+  out_features: *embed_size
+
+similarity_args:
+  embed_size: *embed_size
+  queue_size: 65536
+  temp: 0.07
+
+training_args:
+  log_every_n_steps: 100
+  alpha: 0.4
+  weight_decay: 0.02
+  lr: 1e-5
+  min_lr: 1e-6
+  max_epochs: 5
+  step_size: 100
+  warmup_steps: 1
+  checkpoint_root: "./examples/albef/checkpoints"
+
+eval_args:
+  log_every_n_steps: 100
+  k_test: 256

multimodal/examples/albef/configs/vqa.yaml

@@ -0,0 +1,78 @@
+hidden_size: &hidden_size 768
+vocab_size: &vocab_size 30522
+type_vocab_size: &type_vocab_size 2
+max_position_embeddings: &max_position_embeddings 512
+pad_token_id: &pad_token_id 0
+
+seed: 42
+world_size: 1
+device: "cuda"
+dist_url: "env://"
+output_root: "./examples/albef/outputs"
+
+datamodule_args:
+  train_files: ["./examples/albef/data_files/vqa_train.json", "./examples/albef/data_files/vg_qa.json", "./examples/albef/data_files/vqa_val.json"]
+  test_files: ["./examples/albef/data_files/vqa_test.json"]
+  answer_list: "./examples/albef/data_files/answer_list.json"
+  vqa_root: "./examples/albef/data_files/coco"
+  vg_root: "./examples/albef/data_files/visual_genome"
+  batch_size: 32
+  num_workers: 8
+
+vision_encoder_args:
+  hidden_size: *hidden_size
+  image_size: 384
+  patch_size: 16
+  num_hidden_layers: 12
+  num_attention_heads: 12
+  mlp_dim: 3072
+  dropout: 0.0
+  attention_dropout: 0.0
+  layer_norm_eps: 1e-6
+
+text_encoder_args:
+  vocab_size: *vocab_size
+  hidden_size: *hidden_size
+  type_vocab_size: *type_vocab_size
+  max_position_embeddings: *max_position_embeddings
+  pad_token_id: *pad_token_id
+  num_hidden_layers: 6
+  num_attention_heads: 12
+  intermediate_size: 3072
+  layer_norm_eps: 1e-12
+  dropout: 0.0
+
+multimodal_encoder_args:
+  hidden_size: *hidden_size
+  num_hidden_layers: 6
+  num_attention_heads: 12
+  intermediate_size: 3072
+  layer_norm_eps: 1e-12
+
+text_embeddings_args:
+  hidden_size: *hidden_size
+  vocab_size: *vocab_size
+  pad_token_id: *pad_token_id
+  max_position_embeddings: *max_position_embeddings
+  type_vocab_size: *type_vocab_size
+  layer_norm_eps: 1e-12
+
+prediction_head_args:
+  hidden_size: *hidden_size
+  vocab_size: *vocab_size
+  layer_norm_eps: 1e-12
+
+training_args:
+  log_every_n_steps: 100
+  alpha: 0.4
+  weight_decay: 0.02
+  lr: 2e-5
+  min_lr: 1e-6
+  max_epochs: 8
+  step_size: 100
+  warmup_steps: 4
+  checkpoint_root: "./examples/albef/checkpoints"
+
+eval_args:
+  log_every_n_steps: 100
+  k_test: 128

multimodal/examples/albef/data/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

multimodal/examples/albef/data/retrieval_datamodule.py

@@ -0,0 +1,188 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional, Tuple
+
+import torch
+from data.retrieval_dataset import (
+    ImageToTextRetrievalDataset,
+    RetrievalTrainingDataset,
+    TextToImageRetrievalDataset,
+)
+from data.transforms import (
+    ALBEFTextTransform,
+    testing_image_transform,
+    training_image_transform,
+)
+from pytorch_lightning import LightningDataModule
+from torch import Tensor
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+class RetrievalDataModule(LightningDataModule):
+    """
+    The Data Module for Retrieval task.
+
+    Args:
+        train_files (List[str]): The paths to training json files.
+        test_files (List[str]): The paths to testing json files.
+        image_root (str): The path to image data directory.
+        batch_size (int): The sampling batch size.
+        num_workers (int): The number of workers for the distributed mode.
+    """
+
+    def __init__(
+        self,
+        train_files: List[str],
+        test_files: List[str],
+        image_root: str,
+        batch_size: int,
+        num_workers: int,
+    ) -> None:
+        super().__init__()
+        self.train_dataset = RetrievalTrainingDataset(
+            train_files,
+            image_root,
+            training_image_transform(),
+            ALBEFTextTransform(truncate=True, max_seq_len=30, add_end_token=False),
+        )
+
+        self.image_dataset = ImageToTextRetrievalDataset(
+            test_files,
+            image_root,
+            testing_image_transform(),
+        )
+
+        self.text_dataset = TextToImageRetrievalDataset(
+            test_files,
+            ALBEFTextTransform(
+                truncate=True,
+                pad_to_max_seq_len=True,
+                max_seq_len=30,
+                add_end_token=False,
+            ),
+        )
+
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def _get_sampler(
+        self,
+        dataset: Dataset,
+        shuffle: bool,
+        is_distributed: bool,
+        num_tasks: int,
+        global_rank: int,
+    ) -> Optional[DistributedSampler]:
+        # do not return a sampler if is not in distributed mode
+        # a default RandomSampler is used in this case
+        if not is_distributed:
+            return None
+
+        return DistributedSampler(
+            dataset, num_replicas=num_tasks, rank=global_rank, shuffle=shuffle
+        )
+
+    def train_dataloader(
+        self,
+        is_distributed: bool = False,
+        num_tasks: int = 0,
+        global_rank: int = 0,
+        drop_last: bool = True,
+    ) -> DataLoader:
+        """
+        DataLoader Outputs:
+            images (Tensor): Tensor of shape (B, C, W, H) of image inputs.
+            text (Tensor): Tensor of shape (B, L) of text inputs.
+            text_atts (Tensor): Tensor of shape (B, L) of text attention mask.
+            idx (Tensor): Tensor of shape (B) of image identifiers.
+        """
+        sampler = self._get_sampler(
+            dataset=self.train_dataset,
+            shuffle=True,
+            is_distributed=is_distributed,
+            num_tasks=num_tasks,
+            global_rank=global_rank,
+        )
+        shuffle = sampler is None
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=shuffle,
+            collate_fn=retrieval_train_collate_fn,
+            drop_last=drop_last,
+        )
+
+    def image_dataloader(
+        self,
+        drop_last: bool = False,
+    ) -> DataLoader:
+        """
+        DataLoader Outputs:
+            images (Tensor): Tensor of shape (B, C, W, H) of image inputs.
+        """
+        return DataLoader(
+            self.image_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=None,
+            shuffle=False,
+            collate_fn=None,
+            drop_last=drop_last,
+        )
+
+    def text_dataloader(
+        self,
+        drop_last: bool = False,
+    ) -> DataLoader:
+        """
+        DataLoader Outputs:
+            text (Tensor): Tensor of shape (B, L) of text inputs.
+            text_atts (Tensor): Tensor of shape (B, L) of text attention mask.
+        """
+        return DataLoader(
+            self.text_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=None,
+            shuffle=False,
+            collate_fn=text_collate_fn,
+            drop_last=drop_last,
+        )
+
+
+def retrieval_train_collate_fn(
+    batch: List[Tuple[Tensor, Tensor, int]],
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    image_list = []
+    text_list = []
+    idx_list = []
+    for image, text, idx in batch:
+        image_list.append(image)
+        text_list.append(text)
+        idx_list.append(idx)
+    images = torch.stack(image_list, dim=0)
+    text = pad_sequence(text_list, batch_first=True)
+    text_atts = (text != 0).type(torch.long)
+    idx = Tensor(idx_list).type(torch.long)
+    return (
+        images,
+        text,
+        text_atts,
+        idx,
+    )
+
+
+def text_collate_fn(batch: List[Tensor]) -> Tuple[Tensor, Tensor]:
+    text = pad_sequence(batch, batch_first=True)
+    text_atts = (text != 0).type(torch.long)
+    return text, text_atts

multimodal/examples/albef/data/retrieval_dataset.py

@@ -0,0 +1,149 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+from typing import Callable, List, Tuple, Union
+
+from PIL import Image
+from torch import Tensor
+from torch.utils.data import Dataset
+
+
+class RetrievalTrainingDataset(Dataset):
+    """
+    Create the training dataset for Retrieval task.
+
+    Args:
+        ann_file (List[str]): The paths to training annotation json files.
+        image_root (str): The path to image data directory.
+        image_transform (Callable[[Image.Image], Tensor]): Image data transform.
+        text_transform (Callable[[Union[List[str], str]], Tensor]): Text data transform.
+
+    Dataset Outputs:
+        image (Tensor): Transformed image input tensor of shape (C, H, W).
+        caption (Tensor): Transformed text token input ids.
+        idx (int): The unique identifier for the image.
+    """
+
+    def __init__(
+        self,
+        ann_file: List[str],
+        image_root: str,
+        image_transform: Callable[[Image.Image], Tensor],
+        text_transform: Callable[[Union[List[str], str]], Tensor],
+    ) -> None:
+        self.ann = []
+        for f in ann_file:
+            self.ann += json.load(open(f, "r"))
+
+        self.image_root = image_root
+        self.image_transform = image_transform
+        self.text_transform = text_transform
+
+        self.idx = {}  # map str image_id from dataset to int ids
+        i = 0
+        for ann in self.ann:
+            image_id = ann["image_id"]
+            if image_id not in self.idx.keys():
+                self.idx[image_id] = i
+                i += 1
+
+    def __len__(self) -> int:
+        return len(self.ann)
+
+    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor, int]:
+        ann = self.ann[index]
+        image_path = os.path.join(self.image_root, ann["image"])
+        image = Image.open(image_path).convert("RGB")
+        image = self.image_transform(image)
+        caption = self.text_transform(ann["caption"])
+        return image, caption, self.idx[ann["image_id"]]
+
+
+class ImageToTextRetrievalDataset(Dataset):
+    """
+    Create the dataset for Image-to-Text Retrieval task.
+
+    Args:
+        ann_file (List[str]): The paths to annotation json files.
+        image_root (str): The path to image data directory.
+        image_transform (Callable[[Image.Image], Tensor]): Image data transform.
+
+    Dataset Outputs:
+        image (Tensor): Transformed image input tensor of shape (C, H, W).
+    """
+
+    def __init__(
+        self,
+        ann_file: List[str],
+        image_root: str,
+        image_transform: Callable[[Image.Image], Tensor],
+    ) -> None:
+        self.image_root = image_root
+        self.image_transform = image_transform
+
+        self.ann = []
+        self.images = []  # paths to all images in the dataset
+        self.image_to_text = {}  # map image ids to text ids for evaluation
+        for f in ann_file:
+            self.ann += json.load(open(f, "r"))
+
+        text_id = 0
+        for image_id, ann in enumerate(self.ann):
+            self.images.append(ann["image"])
+            num_text = len(ann["caption"])
+            self.image_to_text[image_id] = list(range(text_id, text_id + num_text))
+            text_id += num_text
+
+    def __len__(self) -> int:
+        return len(self.images)
+
+    def __getitem__(self, index: int) -> Tensor:
+        image_path = os.path.join(self.image_root, self.images[index])
+        image = Image.open(image_path).convert("RGB")
+        image = self.image_transform(image)
+        return image
+
+
+class TextToImageRetrievalDataset(Dataset):
+    """
+    Create the dataset for Text-to-Image Retrieval task.
+
+    Args:
+        ann_file (List[str]): The paths to annotation json files.
+        text_transform (Callable[[Union[List[str], str]], Tensor]): Text data transform.
+
+    Dataset Outputs:
+        text (Tensor): Transformed text token input ids.
+    """
+
+    def __init__(
+        self,
+        ann_file: List[str],
+        text_transform: Callable[[Union[List[str], str]], Tensor],
+    ) -> None:
+        self.text_transform = text_transform
+
+        self.ann = []
+        self.text = []  # all text strings in the dataset
+        self.text_to_image = {}  # map text ids to image ids for evaluation
+        for f in ann_file:
+            self.ann += json.load(open(f, "r"))
+
+        text_id = 0
+        for image_id, ann in enumerate(self.ann):
+            for caption in ann["caption"]:
+                self.text.append(caption)
+                self.text_to_image[text_id] = image_id
+                text_id += 1
+
+    def __len__(self) -> int:
+        return len(self.text)
+
+    def __getitem__(self, index: int) -> Tensor:
+        text = self.text_transform(self.text[index])
+        return text

multimodal/examples/albef/data/transforms.py

@@ -0,0 +1,141 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import re
+from typing import List, Tuple, Union
+
+import torch
+
+from torchtext.transforms import PadTransform, Sequential, ToTensor, Truncate
+from torchvision import transforms
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+# mean and standard deviation from the ALBEF repo:
+# https://github.com/salesforce/ALBEF/blob/main/dataset/__init__.py#L16
+MEAN = (0.48145466, 0.4578275, 0.40821073)
+STD_DEV = (0.26862954, 0.26130258, 0.27577711)
+
+
+class ALBEFTextTransform:
+    """
+    Remove punctuations and trailing spaces in input text and transform it into
+    a Tensor of token ids using BERTTokenizer.
+
+    Args:
+        pretrained_tokenizer (str): Pretrained tokenizer to use.
+            Default: "bert-base-uncased"
+        do_pre_process (bool): Whether to pre-process input text.
+            Defaults to True.
+        truncate (bool): Whether to truncate input text to max_seq_length.
+            Defaults to False.
+        pad_to_max_seq_len (bool): Whether to pad the sequence to max_seq_length.
+        add_end_token (bool): Whether to add the end-of-sentence token.
+            Defaults to True.
+        max_seq_len (int): The max sequence length after truncating or padding.
+            Defaults to 25.
+        cls_token_id (int): Value to represent the start of each text.
+            Defaults to 101, Hugging Face's BERT cls token id.
+        sep_token_id (int): Value to represent the end of each text.
+            Defaults to 102, Hugging Face's BERT sep token id.
+        pad_token_id (int): Value with which to pad each text so that all texts are the same length.
+            Defaults to 0, Hugging Face's BERT pad token id.
+
+    Inputs:
+        text (Union[List[str], str]): Input text to transform.
+    """
+
+    def __init__(
+        self,
+        pretrained_tokenizer: str = "bert-base-uncased",
+        do_pre_process: bool = True,
+        truncate: bool = False,
+        pad_to_max_seq_len: bool = False,
+        add_end_token: bool = True,
+        max_seq_len: int = 25,
+        cls_token_id: int = 101,
+        sep_token_id: int = 102,
+        pad_token_id: int = 0,
+    ):
+        self.do_pre_process = do_pre_process
+        self.cls_token_id = cls_token_id
+        self.sep_token_id = sep_token_id
+        self.pad_token_id = pad_token_id
+        self.add_end_token = add_end_token
+
+        self.tokenizer = BertTokenizer.from_pretrained(pretrained_tokenizer)
+        self.transform = Sequential(
+            Truncate(max_seq_len=max_seq_len) if truncate else torch.nn.Identity(),
+            ToTensor(padding_value=self.pad_token_id),
+            (
+                PadTransform(max_length=max_seq_len, pad_value=self.pad_token_id)
+                if pad_to_max_seq_len
+                else torch.nn.Identity()
+            ),
+        )
+
+    def pre_process(self, text: str) -> str:
+        text = (
+            re.sub(
+                r"([,.'!?\"()*#:;~])",
+                "",
+                text,
+            )
+            .replace("-", " ")
+            .replace("/", " ")
+        )
+        text = text.rstrip(" ")
+
+        return text
+
+    def __call__(self, text: Union[List[str], str]) -> torch.Tensor:
+        if self.do_pre_process:
+            if isinstance(text, str):
+                text = self.pre_process(text)
+            else:
+                text = [self.pre_process(t) for t in text]
+        tokens = self.tokenizer(text)["input_ids"]
+        if not self.add_end_token and tokens[-1] == self.sep_token_id:
+            tokens = tokens[:-1]
+        input_ids = self.transform(tokens)
+
+        return input_ids
+
+
+def training_image_transform(
+    image_size: int = 384,
+    scale: Tuple[float, float] = (0.5, 1.0),
+    image_interpolation=transforms.InterpolationMode.BICUBIC,
+    mean: Tuple[float, float, float] = MEAN,
+    std_dev: Tuple[float, float, float] = STD_DEV,
+) -> transforms.Compose:
+    return transforms.Compose(
+        [
+            transforms.RandomResizedCrop(
+                image_size, scale=scale, interpolation=image_interpolation
+            ),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandAugment(2, 7),
+            transforms.ToTensor(),
+            transforms.Normalize(mean, std_dev),
+        ]
+    )
+
+
+def testing_image_transform(
+    image_size: int = 384,
+    image_interpolation=transforms.InterpolationMode.BICUBIC,
+    mean: Tuple[float, float, float] = MEAN,
+    std_dev: Tuple[float, float, float] = STD_DEV,
+) -> transforms.Compose:
+    return transforms.Compose(
+        [
+            transforms.Resize(
+                (image_size, image_size), interpolation=image_interpolation
+            ),
+            transforms.ToTensor(),
+            transforms.Normalize(mean, std_dev),
+        ]
+    )

multimodal/examples/albef/data/vqa_datamodules.py

@@ -0,0 +1,206 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional, Tuple
+
+import torch
+from data.transforms import (
+    ALBEFTextTransform,
+    testing_image_transform,
+    training_image_transform,
+)
+from data.vqa_dataset import VQADataset
+from pytorch_lightning import LightningDataModule
+from torch import Tensor
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, DistributedSampler
+
+
+class VQADataModule(LightningDataModule):
+    """
+    The Data Module for Visual Question Answering task.
+
+    Args:
+        train_files (List[str]): The paths to training json files.
+        test_files (List[str]): The paths to testing json files.
+        answer_list (str): The path to the answers list.
+        vqa_root (str): The path to vqa data directory.
+        vg_root (str): The path to vg data directory.
+        batch_size (int): The sampling batch size.
+        num_workers (int): The number of workers for the distributed mode.
+    """
+
+    def __init__(
+        self,
+        train_files: List[str],
+        test_files: List[str],
+        answer_list: str,
+        vqa_root: str,
+        vg_root: str,
+        batch_size: int,
+        num_workers: int,
+    ) -> None:
+        super().__init__()
+        self.train_dataset = VQADataset(
+            train_files,
+            vqa_root,
+            vg_root,
+            image_transform=training_image_transform(),
+            question_transform=ALBEFTextTransform(
+                truncate=True, max_seq_len=25, add_end_token=False
+            ),
+            answer_transform=ALBEFTextTransform(do_pre_process=False),
+            split="train",
+        )
+
+        self.test_dataset = VQADataset(
+            test_files,
+            vqa_root,
+            vg_root,
+            image_transform=testing_image_transform(),
+            question_transform=ALBEFTextTransform(add_end_token=False),
+            answer_transform=ALBEFTextTransform(do_pre_process=False),
+            split="test",
+            answer_list=answer_list,
+        )
+
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def _get_sampler(
+        self,
+        dataset: VQADataset,
+        shuffle: bool,
+        is_distributed: bool,
+        num_tasks: int,
+        global_rank: int,
+    ) -> Optional[DistributedSampler]:
+        if not is_distributed:
+            return None
+
+        return DistributedSampler(
+            dataset, num_replicas=num_tasks, rank=global_rank, shuffle=shuffle
+        )
+
+    def train_dataloader(
+        self,
+        is_distributed: bool = False,
+        num_tasks: int = 0,
+        global_rank: int = 0,
+        drop_last: bool = True,
+    ) -> DataLoader:
+        """
+        DataLoader Outputs:
+            images (Tensor): Tensor of shape (B, C, W, H) of image inputs.
+            questions (Tensor): Tensor of shape (B, L) of question inputs.
+            question_atts (Tensor): Tensor of shape (B, L) of question attention mask.
+            answers (Tensor): Tensor of shape (N, M) of answer inputs.
+                N >= B because a vqa sample can have multiple answers.
+            answer_atts (Tensor): Tensor of shape (N, M) of answer attention mask.
+            weights (Tensor): Tensor of shape (N) of answer weights.
+            ans_lengths (List[int]): List of length B and sum N where
+                ans_lengths[i] = number of answers for images[i] and questions[i].
+        """
+        sampler = self._get_sampler(
+            dataset=self.train_dataset,
+            shuffle=True,
+            is_distributed=is_distributed,
+            num_tasks=num_tasks,
+            global_rank=global_rank,
+        )
+        shuffle = sampler is None
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=shuffle,
+            collate_fn=vqa_train_collate_fn,
+            drop_last=drop_last,
+        )
+
+    def test_dataloader(
+        self,
+        is_distributed: bool = False,
+        num_tasks: int = 0,
+        global_rank: int = 0,
+        drop_last=False,
+    ) -> DataLoader:
+        """
+        DataLoader Outputs:
+            images (Tensor): Tensor of shape (B, C, W, H) of image inputs.
+            questions (Tensor): Tensor of shape (B, L) of question inputs.
+            question_atts (Tensor): Tensor of shape (B, L) of question attention mask.
+            question_ids (List): List of length B of question ids.
+        """
+        sampler = self._get_sampler(
+            dataset=self.test_dataset,
+            shuffle=False,
+            is_distributed=is_distributed,
+            num_tasks=num_tasks,
+            global_rank=global_rank,
+        )
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=False,
+            collate_fn=vqa_test_collate_fn,
+            drop_last=drop_last,
+        )
+
+
+def vqa_train_collate_fn(
+    batch: List[Tuple[Tensor, Tensor, List[Tensor], List[float]]],
+) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, List[int]]:
+    image_list = []
+    question_list = []
+    answer_list = []
+    weight_list = []
+    ans_lengths = []
+    for image, question, answer, weights in batch:
+        image_list.append(image)
+        question_list.append(question)
+        answer_list += answer
+        weight_list += weights
+        ans_lengths.append(len(answer))
+    images = torch.stack(image_list, dim=0)
+    questions = pad_sequence(question_list, batch_first=True)
+    question_atts = (questions != 0).type(torch.long)
+    answers = pad_sequence(answer_list, batch_first=True)
+    answer_atts = (answers != 0).type(torch.long)
+    weights = torch.Tensor(weight_list)
+    return (
+        images,
+        questions,
+        question_atts,
+        answers,
+        answer_atts,
+        weights,
+        ans_lengths,
+    )
+
+
+def vqa_test_collate_fn(
+    batch: List[Tuple[Tensor, Tensor, int]],
+) -> Tuple[Tensor, Tensor, Tensor, List[int]]:
+    image_list, question_list, question_ids = [], [], []
+    for image, question, question_id in batch:
+        image_list.append(image)
+        question_list.append(question)
+        question_ids.append(question_id)
+    images = torch.stack(image_list, dim=0)
+    questions = pad_sequence(question_list, batch_first=True)
+    question_atts = (questions != 0).type(torch.long)
+    return (
+        images,
+        questions,
+        question_atts,
+        question_ids,
+    )

multimodal/examples/albef/data/vqa_dataset.py

@@ -0,0 +1,115 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import os
+from typing import Callable, List, Tuple, Union
+
+import torch
+
+from PIL import Image
+from torch import Tensor
+from torch.utils.data import Dataset
+
+
+class VQADataset(Dataset):
+    """
+    Create the dataset for VQA task.
+
+    Args:
+        ann_file (List[str]): The paths to annotation json files.
+        vqa_root (str): The path to vqa data directory.
+        vg_root (str): The path to vg data directory.
+        image_transform (Callable[[Image.Image], Tensor]): image data transform.
+        question_transform (Callable[[Union[List[str], str]], Tensor]): text data transform for questions.
+        answer_transform (Callable[[Union[List[str], str]], Tensor]): text data transform for answers.
+        split (str): Indicates train or test. Default is train.
+        answer_list (str): The path to the answers list. Required for test split.
+
+    Dataset Outputs:
+        if split is train:
+            image (Tensor): Transformed image input tensor of shape (C, W, H).
+            question (Tensor): Transformed question token input ids.
+            answers (List[Tensor]): List of transformed answers token input ids.
+            answer_weights (List[float]): List of answer weights.
+                answer_weights[i] is proportional to the number of occurences of answers[i]
+        if split is test:
+            image (Tensor): Transformed image input tensor of shape (C, W, H).
+            question (Tensor): Transformed text token input ids.
+            question_id (int): The question sample id.
+    """
+
+    def __init__(
+        self,
+        ann_file: List[str],
+        vqa_root: str,
+        vg_root: str,
+        image_transform: Callable[[Image.Image], Tensor],
+        question_transform: Callable[[Union[List[str], str]], Tensor],
+        answer_transform: Callable[[Union[List[str], str]], Tensor],
+        split: str = "train",
+        answer_list: str = None,
+    ) -> None:
+        self.ann = []
+        for f in ann_file:
+            self.ann += json.load(open(f, "r"))
+
+        self.vqa_root = vqa_root
+        self.vg_root = vg_root
+        self.image_transform = image_transform
+        self.question_transform = question_transform
+        self.answer_transform = answer_transform
+        self.split = split
+
+        if split == "test":
+            self.answer_list = json.load(open(answer_list, "r"))
+            self.answer_input_ids = self.answer_transform(self.answer_list)
+            self.answer_attention_mask = (self.answer_input_ids != 0).type(torch.long)
+
+    def __len__(self) -> int:
+        return len(self.ann)
+
+    def __getitem__(
+        self, index: int
+    ) -> Union[
+        Tuple[Tensor, Tensor, int], Tuple[Tensor, Tensor, List[Tensor], List[float]]
+    ]:
+        ann = self.ann[index]
+
+        image_root = self.vqa_root if ann["dataset"] == "vqa" else self.vg_root
+        image_path = os.path.join(image_root, ann["image"])
+        image = Image.open(image_path).convert("RGB")
+        image = self.image_transform(image)
+        question = self.question_transform(ann["question"])
+
+        if self.split == "test":
+            return image, question, ann["question_id"]
+
+        elif self.split == "train":
+            if ann["dataset"] == "vqa":
+                # Each VQA sample question has a list of answers (with potential repeats)
+                # answer_weight[answer] = count(answer) / len(answers for the question)
+                answer_weights = {}
+                for answer in ann["answer"]:
+                    if answer in answer_weights.keys():
+                        answer_weights[answer] += 1 / len(ann["answer"])
+                    else:
+                        answer_weights[answer] = 1 / len(ann["answer"])
+
+                answers = list(answer_weights.keys())
+                answer_weights = list(answer_weights.values())
+
+            elif ann["dataset"] == "vg":
+                # A VG sample question has one answer so assign it a constant weight (0.5)
+                answers = [ann["answer"]]
+                answer_weights = [0.5]
+
+            answers = list(self.answer_transform(answers))
+
+            return image, question, answers, answer_weights
+
+        else:
+            raise ValueError("dataset split should be train or test")

multimodal/examples/common/data/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .multidata import *  # noqa F401

multimodal/examples/common/data/multidata.py

@@ -0,0 +1,194 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+import warnings
+from functools import partial
+from typing import Callable, List, Optional
+
+import torch
+from pytorch_lightning import LightningDataModule
+
+
+class MultiDataLoader:
+    # NOTE: Please check MMF's MultiDataLoader if you want to support
+    # epoch based sampling funcs.
+    def __init__(
+        self,
+        loaders: List[torch.utils.data.DataLoader],
+        sampling_func: Optional[Callable] = None,
+    ):
+        """MultiDataLoader takes in a list of dataloaders and a sampling function
+        and cycles between these dataloaders after each batch based on the index
+        provided by the sampling function passed. Useful for doing multi-tasking
+        over multiple datasets
+
+        Args:
+            loaders (List[torch.utils.data.DataLoader]): List of dataloaders on
+                which the multitasking has to be done.
+
+            sampling_func (Optional[Callable], optional): Function which will return
+                the next index to be selected. Defaults to equally weight sampling.
+        """
+        if loaders is None or len(loaders) == 0:
+            warnings.warn(
+                "Empty loaders passed into MultiDataLoader. This can have "
+                "unintended consequences."
+            )
+
+        if sampling_func is None:
+            sampling_func = partial(random.choice, range(len(loaders)))
+
+        self.sampling_func = sampling_func
+        self.loaders = loaders
+        self.num_datasets = len(self.loaders)
+        self.iterators = [None for _ in loaders]
+        self.current_index = 0
+        self.set_samplers()
+
+    def set_samplers(self):
+        self.samplers: List[torch.utils.data.Sampler] = []
+        for loader in self.loaders:
+            if hasattr(loader, "sampler"):
+                self.samplers.append(loader.sampler)
+
+    def __iter__(self):
+        self.iterators = []
+
+        for loader in self.loaders:
+            self.iterators.append(iter(loader))
+
+        self.change_dataloader()
+
+        return self
+
+    def __next__(self):
+        """
+        Calculation of next batch is performed using following logic.
+
+        Current chosen iterator is set in the change_dataloader function
+        based on the `sampling_func` function passed to `__init__` of the
+        dataloader which is called to get the index of next selected dataloader.
+
+        If we get the next batch from iterator without any StopIteration exception,
+        we return it as it is.
+
+        Epochs don't make sense in case of using `sampling_func` unless you add
+        extra logic to support epoch-based sampling functions. MMF does this in
+        a different way, so take a look at IterationStrategies there to understand
+        how this can be possibly done.
+
+        Think of a case of random (equal) proportional sampling for dataset x and y
+        where x is half the size of y. When x will complete its 2 epochs, y will
+        have only 1 epoch completed. **So please don't use max_epochs or epoch
+        based training in this case as it won't be honored**. If an iterator is
+        finished, we just reignite it in this case and finished iterators
+        variable isn't used. This means that this case will never reach the
+        __iter__ function ever again.
+
+
+        Returns:
+            Dict: Contains two keys, one "batch" containing the batch from current
+                selected dataloader and "datamodule_index" which is index of
+                currently selected dataloader.
+        """
+        self.change_dataloader()
+        try:
+            next_batch = next(self.current_iterator)
+        except StopIteration:
+            iterator = iter(self.loaders[self.current_index])
+            self.iterators[self.current_index] = iterator
+            self.current_iterator = iterator
+            next_batch = next(self.current_iterator)
+
+        return {"batch": next_batch, "datamodule_index": self.current_index}
+
+    def change_dataloader(self):
+        choice = 0
+
+        if self.num_datasets <= 1:
+            self.current_index = choice
+            self.current_iterator = self.iterators[self.current_index]
+            return
+
+        choice = [self.sampling_func()]
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            # This broadcast is probably unnecessary with lightning if everything
+            # is already properly seeded. But,to be on safe side, we can still
+            # do this.
+            # There are also some smarter ways to do this to avoid any broadcasting
+            # by basically having a fixed generator with a fixed seed which will
+            # always work deterministically.
+            # TODO: Check if not doing this provides any speed benefits.
+            torch.distributed.broadcast_object_list(choice, 0)
+
+        self.current_index = choice[0]
+        self.current_iterator = self.iterators[self.current_index]
+
+    def set_epoch(self, epoch: int):
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            for sampler in self.samplers:
+                if sampler is not None and hasattr(sampler, "set_epoch"):
+                    sampler.set_epoch(epoch)
+
+
+class MultiDataModule(LightningDataModule):
+    """MultiDataModule is just an abstraction over MultiDataLoader
+    that will allow us to integrate it with Lightning.
+    """
+
+    # NOTE: Add rest of the functions that should be called on child datamodules
+    # as required
+    def __init__(
+        self,
+        datamodules: List[LightningDataModule],
+        sampling_func: Optional[Callable] = None,
+    ):
+        super().__init__()
+        self.datamodules = datamodules
+        self.sampling_func = sampling_func
+        self.current_datamodule_idx = 0
+
+    def setup(self, stage=None):
+        for datamodule in self.datamodules:
+            datamodule.setup(stage)
+
+    def prepare_data(self):
+        for datamodule in self.datamodules:
+            datamodule.prepare_data()
+
+    def train_dataloader(self) -> MultiDataLoader:
+        # TODO: Fix assign inconsistency
+        return self._build_multi_dataloader("train")
+
+    def val_dataloader(self) -> MultiDataLoader:
+        return self._build_multi_dataloader("val")
+
+    def test_dataloader(self) -> MultiDataLoader:
+        return self._build_multi_dataloader("test")
+
+    def _build_multi_dataloader(self, split="train"):
+        dataloaders = []
+        for datamodule in self.datamodules:
+            dataloaders.append(getattr(datamodule, f"{split}_dataloader")())
+
+        return MultiDataLoader(dataloaders, self.sampling_func)
+
+    def on_before_batch_transfer(self, batch, *args):
+        batch, index = batch["batch"], batch["datamodule_index"]
+        self.current_datamodule_idx = index
+        return self.datamodules[self.current_datamodule_idx].on_before_batch_transfer(
+            batch, *args
+        )
+
+    def on_after_batch_transfer(self, batch, *args):
+        return self.datamodules[self.current_datamodule_idx].on_after_batch_transfer(
+            batch, *args
+        )
+
+    def teardown(self, stage):
+        for datamodule in self.datamodules:
+            datamodule.teardown(stage)

multimodal/examples/flava/callbacks/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .multimodal_eval import *  # noqa F401

multimodal/examples/flava/callbacks/multimodal_eval.py

@@ -0,0 +1,108 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import logging
+
+import torch
+from flava.data import default_text_transform, VL_MAX_LENGTH_DEFAULT
+from flava.data.imagenet_zeroshot_data import (
+    imagenet_classnames,
+    openai_imagenet_template,
+)
+from pytorch_lightning import Callback, LightningDataModule
+from pytorch_lightning.utilities import rank_zero_only
+from tqdm import tqdm
+
+
+logger = logging.getLogger(__name__)
+
+
+def _zero_shot_classifier(model, device, text_transform, *args, **kwargs):
+    zeroshot_weights = []
+    for classname in tqdm(imagenet_classnames):
+        texts = text_transform(
+            [template(classname) for template in openai_imagenet_template]
+        )["input_ids"]
+        texts = texts.to(device)
+        class_embeddings = model.encode_text(texts)
+        class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+        class_embedding = class_embeddings.mean(dim=0)
+        class_embedding /= class_embedding.norm()
+        zeroshot_weights.append(class_embedding)
+
+    zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
+    return zeroshot_weights
+
+
+def _accuracy(output, target, topk=(1,)):
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [
+        float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy())
+        for k in topk
+    ]
+
+
+@rank_zero_only
+def run_imagenet_zero_shot(model, dataloader, device, text_transform, *args, **kwargs):
+    logger.info("Starting ImageNet Zero-Shot Eval")
+    logger.info("Building classifier")
+    classifier = _zero_shot_classifier(model, device, text_transform)
+    logger.info("Classifier built")
+    top1, top5, n = 0.0, 0.0, 0.0
+    for sample in tqdm(dataloader):
+        images = sample["image"]
+        target = sample["label"]
+        images = images.to(device)
+        target = target.to(device)
+
+        # predict
+        # if hasattr(model, "module"):
+        #     image_features = model.module.encode_image({"image": images})
+        # else:
+        image_features = model.encode_image(images)
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        logits = 100.0 * image_features @ classifier
+
+        # measure accuracy
+        acc1, acc5 = _accuracy(logits, target, topk=(1, 5))
+        top1 += acc1
+        top5 += acc5
+        n += images.size(0)
+
+    top1 = top1 / n
+    top5 = top5 / n
+    results = {}
+    results["imagenet-zeroshot-val-top1"] = top1
+    results["imagenet-zeroshot-val-top5"] = top5
+    return results
+
+
+class MultimodalEvalCallback(Callback):
+    def __init__(self, imagenet_datamodule: LightningDataModule, *args, **kwargs):
+        super().__init__()
+        self.imagenet_val_dataloader = imagenet_datamodule.val_dataloader()
+        self.text_transform = default_text_transform(
+            max_text_length=VL_MAX_LENGTH_DEFAULT
+        )
+
+    @torch.no_grad()
+    def on_validation_start(self, trainer, pl_module, **kwargs) -> None:
+        metrics = run_imagenet_zero_shot(
+            pl_module.model,
+            self.imagenet_val_dataloader,
+            pl_module.device,
+            self.text_transform,
+        )
+        if metrics is not None:
+            for key in metrics:
+                self.log(
+                    f"val/{key}",
+                    metrics[key],
+                    prog_bar=True,
+                    logger=True,
+                    rank_zero_only=True,
+                )

multimodal/examples/flava/configs/finetuning/qnli.yaml

@@ -0,0 +1,48 @@
+# Note that in original FLAVA paper, only Logistic Regression numbers were provided for image datasets.
+_target_: flava.definitions.FLAVAArguments
+training:
+  _target_: flava.definitions.TrainingArguments
+  lightning:
+    max_steps: 33112
+    gpus: 1
+    val_check_interval: 1000
+    num_sanity_val_steps: 0
+    strategy: ddp
+  lightning_checkpoint:
+    dirpath: "."
+    filename: flava-{epoch:02d}-{step}
+    save_last: true
+    every_n_train_steps: 1000
+    save_on_train_epoch_end: true
+    verbose: true
+    monitor: validation/accuracy/classification
+    mode: max
+  lightning_load_from_checkpoint: null
+  seed: -1
+  batch_size: 32
+  num_workers: 4
+  learning_rate: 1e-5
+  adam_eps: 1e-6
+  adam_weight_decay: 0.1
+  adam_betas:
+  - 0.9
+  - 0.98
+  warmup_steps: 1986
+
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - text
+  num_classes: 2
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: glue
+        subset: qnli
+        rename_columns:
+          - ["question", "sentence1"]
+          - ["sentence", "sentence2"]
+    datamodule_extra_kwargs:
+      text_columns: ["sentence1", "sentence2"]

multimodal/examples/flava/configs/finetuning/rendered_sst2.yaml

@@ -0,0 +1,37 @@
+# Note that in original FLAVA paper, only Logistic Regression numbers were provided for image datasets.
+_target_: flava.definitions.FLAVAArguments
+training:
+  _target_: flava.definitions.TrainingArguments
+  lightning:
+    max_steps: 20935
+    gpus: -1
+    val_check_interval: 100
+    num_sanity_val_steps: 0
+    strategy: ddp
+  lightning_checkpoint:
+    dirpath: "."
+    filename: flava-{epoch:02d}-{step}
+    save_last: true
+    every_n_train_steps: 1000
+    save_on_train_epoch_end: true
+    verbose: true
+  lightning_load_from_checkpoint: null
+  seed: -1
+  batch_size: 32
+  num_workers: 4
+  learning_rate: 1e-5
+  adam_eps: 1e-8
+  adam_weight_decay: 1e-2
+  warmup_steps: 1256
+
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  num_classes: 2
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.TorchVisionDatasetInfo
+        key: RenderedSST2

multimodal/examples/flava/configs/pretraining/debug.yaml

@@ -0,0 +1,61 @@
+_target_: flava.definitions.FLAVAArguments
+training:
+  _target_: flava.definitions.TrainingArguments
+  lightning:
+    max_steps: 450000
+    gpus: -1
+    val_check_interval: 10000
+    num_sanity_val_steps: 0
+    strategy: ddp
+  lightning_checkpoint:
+    dirpath: "."
+    filename: flava-{epoch:02d}-{step}
+    save_last: true
+    every_n_train_steps: 1000
+    save_on_train_epoch_end: true
+    verbose: true
+  lightning_load_from_checkpoint: null
+  seed: -1
+  batch_size: 8
+  num_workers: 4
+  learning_rate: 2e-4
+  adam_eps: 1e-8
+  adam_weight_decay: 1e-2
+  warmup_steps: 2000
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: jellyfish
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: jellyfish
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train

multimodal/examples/flava/data/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .transforms import *  # noqa F401
+from .utils import *  # noqa F401
+from .imagenet_zeroshot_data import *  # noqa F401
+from .datamodules import *  # noqa F401

multimodal/examples/flava/data/datamodules.py

@@ -0,0 +1,529 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torchvision
+from flava.definitions import HFDatasetInfo, TorchVisionDatasetInfo
+from pytorch_lightning import LightningDataModule
+from transformers import (
+    BertTokenizer,
+    DataCollatorForLanguageModeling,
+    DataCollatorForWholeWordMask,
+    DefaultDataCollator,
+    TRANSFORMERS_CACHE,
+)
+from transformers.data.data_collator import torch_default_data_collator
+
+from .transforms import (
+    default_image_pretraining_transforms,
+    default_text_transform,
+    default_torchvision_transforms,
+    encode_text_batch,
+    pad_batch,
+    TEXT_DEFAULT_TOKENIZER,
+    TEXT_WHOLE_WORD_MASK_TOKENIZER,
+    VL_MAX_LENGTH_DEFAULT,
+    VLTransform,
+)
+from .utils import build_datasets_from_info, fetch_images
+
+
+def transform_image(transform, sample):
+    sample.update(transform(sample["image"]))
+    return sample
+
+
+class DataCollatorForWholeWordMaskRetainingBatch(DataCollatorForWholeWordMask):
+    def torch_call(
+        self, examples: List[Union[List[int], Any, Dict[str, Any]]]
+    ) -> Dict[str, Any]:
+        masked_batch = super().torch_call(examples)
+        examples = torch_default_data_collator(examples)
+        examples["input_ids"] = masked_batch["input_ids"]
+        examples["labels"] = masked_batch["labels"]
+        return examples
+
+
+class ImageDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        transforms: Optional[Tuple[Callable, Callable]] = None,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        allow_uneven_batches: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_dataset_infos = train_infos
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.allow_uneven_batches = allow_uneven_batches
+
+        if transforms is None:
+            transforms = default_image_pretraining_transforms()
+
+        self.train_transform, self.test_transform = transforms
+
+    def setup(self, stage=None):
+        train_transform = partial(transform_image, self.train_transform)
+        val_transform = partial(transform_image, self.test_transform)
+
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(train_transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(val_transform)
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=None,
+            shuffle=True,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            # ideally, we don't need to drop these for unimodal cases
+            # but just to be safe
+            drop_last=True,
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=None,
+            shuffle=False,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            # ideally, we don't need to drop these for unimodal cases
+            # but just to be safe
+            drop_last=True,
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()
+
+    def on_before_batch_transfer(self, batch, *args):
+        if batch["label"].size(0) < self.batch_size and not self.allow_uneven_batches:
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+
+class TextDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        text_columns: List[str],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        tokenizer: Optional[Callable] = None,
+        max_length: int = 512,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        allow_uneven_batches: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_dataset_infos = train_infos
+        self.text_columns = text_columns
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.allow_uneven_batches = allow_uneven_batches
+
+    def setup(self, stage=None):
+        if self.tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained(TEXT_DEFAULT_TOKENIZER)
+        transform = partial(
+            encode_text_batch,
+            tokenizer=self.tokenizer,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_tensors="pt",
+            return_special_tokens_mask=True,
+            text_columns=self.text_columns,
+            return_batch=True,
+        )
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(transform)
+
+    def train_dataloader(self):
+        return self._build_dataloader(self.train_dataset)
+
+    def val_dataloader(self):
+        return self._build_dataloader(self.val_dataset, shuffle=False)
+
+    def _build_dataloader(self, dataset, drop_last=False, shuffle=True):
+        return torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=None,
+            shuffle=shuffle,
+            collate_fn=self._build_collator(),
+            drop_last=drop_last,
+        )
+
+    def _build_collator(self):
+        return DefaultDataCollator()
+
+    def on_before_batch_transfer(self, batch, *args):
+        batch.pop("token_type_ids", None)
+        mask = batch.pop("attention_mask", None)
+        if mask.size(0) < self.batch_size and not self.allow_uneven_batches:
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+    def on_after_batch_transfer(self, batch, *args):
+        batch["text"] = batch.pop("input_ids")
+        return batch
+
+
+class MLMDataModule(TextDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        text_columns: List[str],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        mlm_probability: float = 0.15,
+        ignore_index: int = -1,
+        **kwargs: Any,
+    ):
+        super().__init__(train_infos, text_columns, val_infos, **kwargs)
+        self.mlm_probability = mlm_probability
+        self.ignore_index = ignore_index
+
+    def setup(self, stage=None):
+        if self.tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained(TEXT_DEFAULT_TOKENIZER)
+        transform = partial(
+            encode_text_batch,
+            tokenizer=self.tokenizer,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_tensors="pt",
+            return_special_tokens_mask=True,
+            text_columns=self.text_columns,
+            return_batch=False,
+        )
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(transform)
+
+    def _build_dataloader(self, dataset, drop_last=True, shuffle=True):
+        # uneven batches can cause distributed issues,
+        # drop last batch to prevent those.
+        # ideally, we don't need to drop these for unimodal cases
+        # but just to be safe
+        return super()._build_dataloader(dataset, drop_last=drop_last, shuffle=shuffle)
+
+    def _build_collator(self):
+        return DataCollatorForLanguageModeling(
+            self.tokenizer, mlm_probability=self.mlm_probability
+        )
+
+    def on_after_batch_transfer(self, batch, *args):
+        batch["text_masked"] = batch.pop("input_ids")
+        batch["mlm_labels"] = batch.pop("labels")
+        batch["mlm_labels"][batch["mlm_labels"] == -100] = self.ignore_index
+        return batch
+
+
+class VLDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        val_infos: List[HFDatasetInfo],
+        text_transform: Optional[Callable] = None,
+        image_transforms: Optional[Tuple[Callable, Callable]] = None,
+        mlm_probablity: float = 0.15,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        finetuning: bool = False,
+        ignore_index: int = -1,
+        itm_probability: float = 0.1,
+        allow_uneven_batches: bool = False,
+        fetch_num_threads: int = 4,
+        fetch_retries: int = 0,
+        fetch_sleep_timer: int = 0,
+        fetch_timeout: Optional[float] = None,
+        fetch_batch_size: int = 50,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.train_dataset_infos = train_infos
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+        if image_transforms is None:
+            if not finetuning:
+                image_transforms = default_image_pretraining_transforms()
+            else:
+                image_transforms = default_torchvision_transforms(use_dict=True)
+
+        self.train_image_transform, self.test_image_transform = image_transforms
+        self.text_transform = text_transform
+        self.mlm_probability = mlm_probablity
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.ignore_index = ignore_index
+        self.itm_probability = itm_probability
+        self.allow_uneven_batches = allow_uneven_batches
+        self.fetch_num_threads = fetch_num_threads
+        self.fetch_retries = fetch_retries
+        self.fetch_sleep_timer = fetch_sleep_timer
+        self.fetch_timeout = fetch_timeout
+        self.fetch_batch_size = fetch_batch_size
+
+    def setup(self, stage=None):
+        if self.text_transform is None:
+            # TODO Update to use whole word mask vocab
+            text_tokenizer = BertTokenizer.from_pretrained(
+                TEXT_WHOLE_WORD_MASK_TOKENIZER
+            )
+            self.text_transform = default_text_transform(
+                text_tokenizer, max_text_length=VL_MAX_LENGTH_DEFAULT
+            )
+        self.text_tokenizer = self.text_transform.keywords["tokenizer"]
+        train_vl_transform = VLTransform(
+            self.train_image_transform, self.text_transform
+        )
+        val_vl_transform = VLTransform(self.test_image_transform, self.text_transform)
+
+        train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        train_dataset = train_dataset.map(
+            fetch_images,
+            batched=True,
+            batch_size=self.fetch_batch_size,
+            fn_kwargs={
+                "num_threads": self.fetch_num_threads,
+                "timeout": self.fetch_timeout,
+                "retries": self.fetch_retries,
+                "sleep_timer": self.fetch_sleep_timer,
+            },
+        )
+        train_dataset = train_dataset.filter(
+            lambda example: example["image"] is not None
+        )
+        self.train_dataset = train_dataset
+        self.train_dataset.set_transform(
+            partial(
+                train_vl_transform,
+                dataset=train_dataset.filter(lambda example: True),
+                itm_probability=self.itm_probability,
+            )
+        )
+
+        val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+
+        val_dataset = val_dataset.map(
+            fetch_images,
+            batched=True,
+            batch_size=self.fetch_batch_size,
+            fn_kwargs={
+                "num_threads": self.fetch_num_threads,
+                "timeout": self.fetch_timeout,
+                "retries": self.fetch_retries,
+                "sleep_timer": self.fetch_sleep_timer,
+            },
+        )
+        val_dataset = val_dataset.filter(lambda example: example["image"] is not None)
+        self.val_dataset = val_dataset
+        self.val_dataset.set_transform(
+            partial(
+                val_vl_transform,
+                dataset=self.val_dataset.filter(
+                    lambda example: True
+                ),  # Pass a copy to transform
+                itm_probability=self.itm_probability,
+            )
+        )
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=None,
+            shuffle=True,
+            collate_fn=self._build_collator(),
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            drop_last=True,
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=None,
+            shuffle=False,
+            collate_fn=self._build_collator(),
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            drop_last=True,
+        )
+
+    def _build_collator(self):
+        return DataCollatorForWholeWordMaskRetainingBatch(
+            self.text_tokenizer, mlm_probability=self.mlm_probability
+        )
+
+    def on_before_batch_transfer(self, batch, *args):
+        batch.pop("token_type_ids", None)
+        mask = batch.pop("attention_mask", None)
+        if (
+            mask is not None
+            and mask.size(0) < self.batch_size
+            and not self.allow_uneven_batches
+        ):
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+    def on_after_batch_transfer(self, batch, *args):
+        text_masked = batch.pop("input_ids")
+        mlm_labels = batch.pop("labels", None)
+        mlm_labels[mlm_labels == -100] = self.ignore_index
+        text = text_masked.detach().clone()
+        text[mlm_labels != -1] = mlm_labels[mlm_labels != -1]
+        batch.update(
+            {"mlm_labels": mlm_labels, "text": text, "text_masked": text_masked}
+        )
+        return batch
+
+
+class TorchVisionDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[TorchVisionDatasetInfo],
+        # Val info is not used for torchvision datamodule, but kept to keep things consistent
+        val_infos: Optional[List[TorchVisionDatasetInfo]] = None,
+        dataset_root: Optional[str] = None,
+        image_transforms: Optional[Tuple[Callable, Callable]] = None,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_info = train_infos[0]
+        if val_infos is None:
+            val_infos = train_infos
+        self.val_info = val_infos[0]
+
+        self.train_class_ptr, self.train_root = self._parse_info(
+            self.train_info, dataset_root=dataset_root
+        )
+        self.val_class_ptr, self.val_root = self._parse_info(
+            self.val_info, dataset_root=dataset_root
+        )
+
+        if image_transforms is None:
+            image_transforms = default_torchvision_transforms()
+
+        self.train_transform, self.test_transform = image_transforms
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+
+    def _parse_info(
+        self, info: TorchVisionDatasetInfo, dataset_root: Optional[str] = None
+    ):
+        assert hasattr(
+            torchvision.datasets, info.key
+        ), f"No dataset named {info.key} present in torchvision.datasets"
+        class_ptr = getattr(torchvision.datasets, info.key)
+        if dataset_root is None:
+            dataset_root = os.path.join(TRANSFORMERS_CACHE, "datasets", "torchvision")
+            dataset_root = os.path.join(dataset_root, class_ptr.__name__.lower())
+            os.makedirs(dataset_root, exist_ok=True)
+
+        return class_ptr, dataset_root
+
+    def setup(self, stage=None):
+        self.train_dataset = self.train_class_ptr(
+            self.train_root,
+            split=self.train_info.train_split,
+            transform=self.train_transform,
+            download=True,
+        )
+
+        if self.val_info.has_val:
+            self.val_dataset = self.val_class_ptr(
+                self.val_root,
+                split=self.val_info.val_split,
+                transform=self.test_transform,
+                download=True,
+            )
+
+        self.test_dataset = self.val_class_ptr(
+            self.val_root,
+            split=self.val_info.test_split,
+            transform=self.test_transform,
+            download=True,
+        )
+
+    def train_dataloader(self):
+        return self._build_dataloader(self.train_dataset)
+
+    def val_dataloader(self):
+        if self.val_info.has_val:
+            dataset = self.val_dataset
+        else:
+            dataset = self.test_dataset
+
+        return self._build_dataloader(dataset, shuffle=False)
+
+    def test_dataloader(self):
+        return self._build_dataloader(self.test_dataset, shuffle=False)
+
+    def _build_dataloader(self, dataset: torch.utils.data.Dataset, shuffle=True):
+        return torch.utils.data.DataLoader(
+            dataset,
+            shuffle=shuffle,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+        )
+
+    def on_before_batch_transfer(self, batch, *args):
+        images, targets = batch
+        batch = {"image": images, "labels": targets}
+        return batch

multimodal/examples/flava/data/imagenet_zeroshot_data.py

@@ -0,0 +1,1095 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# File taken from https://github.com/mlfoundations/open_clip/
+
+
+imagenet_classnames = [
+    "tench",
+    "goldfish",
+    "great white shark",
+    "tiger shark",
+    "hammerhead shark",
+    "electric ray",
+    "stingray",
+    "rooster",
+    "hen",
+    "ostrich",
+    "brambling",
+    "goldfinch",
+    "house finch",
+    "junco",
+    "indigo bunting",
+    "American robin",
+    "bulbul",
+    "jay",
+    "magpie",
+    "chickadee",
+    "American dipper",
+    "kite (bird of prey)",
+    "bald eagle",
+    "vulture",
+    "great grey owl",
+    "fire salamander",
+    "smooth newt",
+    "newt",
+    "spotted salamander",
+    "axolotl",
+    "American bullfrog",
+    "tree frog",
+    "tailed frog",
+    "loggerhead sea turtle",
+    "leatherback sea turtle",
+    "mud turtle",
+    "terrapin",
+    "box turtle",
+    "banded gecko",
+    "green iguana",
+    "Carolina anole",
+    "desert grassland whiptail lizard",
+    "agama",
+    "frilled-necked lizard",
+    "alligator lizard",
+    "Gila monster",
+    "European green lizard",
+    "chameleon",
+    "Komodo dragon",
+    "Nile crocodile",
+    "American alligator",
+    "triceratops",
+    "worm snake",
+    "ring-necked snake",
+    "eastern hog-nosed snake",
+    "smooth green snake",
+    "kingsnake",
+    "garter snake",
+    "water snake",
+    "vine snake",
+    "night snake",
+    "boa constrictor",
+    "African rock python",
+    "Indian cobra",
+    "green mamba",
+    "sea snake",
+    "Saharan horned viper",
+    "eastern diamondback rattlesnake",
+    "sidewinder rattlesnake",
+    "trilobite",
+    "harvestman",
+    "scorpion",
+    "yellow garden spider",
+    "barn spider",
+    "European garden spider",
+    "southern black widow",
+    "tarantula",
+    "wolf spider",
+    "tick",
+    "centipede",
+    "black grouse",
+    "ptarmigan",
+    "ruffed grouse",
+    "prairie grouse",
+    "peafowl",
+    "quail",
+    "partridge",
+    "african grey parrot",
+    "macaw",
+    "sulphur-crested cockatoo",
+    "lorikeet",
+    "coucal",
+    "bee eater",
+    "hornbill",
+    "hummingbird",
+    "jacamar",
+    "toucan",
+    "duck",
+    "red-breasted merganser",
+    "goose",
+    "black swan",
+    "tusker",
+    "echidna",
+    "platypus",
+    "wallaby",
+    "koala",
+    "wombat",
+    "jellyfish",
+    "sea anemone",
+    "brain coral",
+    "flatworm",
+    "nematode",
+    "conch",
+    "snail",
+    "slug",
+    "sea slug",
+    "chiton",
+    "chambered nautilus",
+    "Dungeness crab",
+    "rock crab",
+    "fiddler crab",
+    "red king crab",
+    "American lobster",
+    "spiny lobster",
+    "crayfish",
+    "hermit crab",
+    "isopod",
+    "white stork",
+    "black stork",
+    "spoonbill",
+    "flamingo",
+    "little blue heron",
+    "great egret",
+    "bittern bird",
+    "crane bird",
+    "limpkin",
+    "common gallinule",
+    "American coot",
+    "bustard",
+    "ruddy turnstone",
+    "dunlin",
+    "common redshank",
+    "dowitcher",
+    "oystercatcher",
+    "pelican",
+    "king penguin",
+    "albatross",
+    "grey whale",
+    "killer whale",
+    "dugong",
+    "sea lion",
+    "Chihuahua",
+    "Japanese Chin",
+    "Maltese",
+    "Pekingese",
+    "Shih Tzu",
+    "King Charles Spaniel",
+    "Papillon",
+    "toy terrier",
+    "Rhodesian Ridgeback",
+    "Afghan Hound",
+    "Basset Hound",
+    "Beagle",
+    "Bloodhound",
+    "Bluetick Coonhound",
+    "Black and Tan Coonhound",
+    "Treeing Walker Coonhound",
+    "English foxhound",
+    "Redbone Coonhound",
+    "borzoi",
+    "Irish Wolfhound",
+    "Italian Greyhound",
+    "Whippet",
+    "Ibizan Hound",
+    "Norwegian Elkhound",
+    "Otterhound",
+    "Saluki",
+    "Scottish Deerhound",
+    "Weimaraner",
+    "Staffordshire Bull Terrier",
+    "American Staffordshire Terrier",
+    "Bedlington Terrier",
+    "Border Terrier",
+    "Kerry Blue Terrier",
+    "Irish Terrier",
+    "Norfolk Terrier",
+    "Norwich Terrier",
+    "Yorkshire Terrier",
+    "Wire Fox Terrier",
+    "Lakeland Terrier",
+    "Sealyham Terrier",
+    "Airedale Terrier",
+    "Cairn Terrier",
+    "Australian Terrier",
+    "Dandie Dinmont Terrier",
+    "Boston Terrier",
+    "Miniature Schnauzer",
+    "Giant Schnauzer",
+    "Standard Schnauzer",
+    "Scottish Terrier",
+    "Tibetan Terrier",
+    "Australian Silky Terrier",
+    "Soft-coated Wheaten Terrier",
+    "West Highland White Terrier",
+    "Lhasa Apso",
+    "Flat-Coated Retriever",
+    "Curly-coated Retriever",
+    "Golden Retriever",
+    "Labrador Retriever",
+    "Chesapeake Bay Retriever",
+    "German Shorthaired Pointer",
+    "Vizsla",
+    "English Setter",
+    "Irish Setter",
+    "Gordon Setter",
+    "Brittany dog",
+    "Clumber Spaniel",
+    "English Springer Spaniel",
+    "Welsh Springer Spaniel",
+    "Cocker Spaniel",
+    "Sussex Spaniel",
+    "Irish Water Spaniel",
+    "Kuvasz",
+    "Schipperke",
+    "Groenendael dog",
+    "Malinois",
+    "Briard",
+    "Australian Kelpie",
+    "Komondor",
+    "Old English Sheepdog",
+    "Shetland Sheepdog",
+    "collie",
+    "Border Collie",
+    "Bouvier des Flandres dog",
+    "Rottweiler",
+    "German Shepherd Dog",
+    "Dobermann",
+    "Miniature Pinscher",
+    "Greater Swiss Mountain Dog",
+    "Bernese Mountain Dog",
+    "Appenzeller Sennenhund",
+    "Entlebucher Sennenhund",
+    "Boxer",
+    "Bullmastiff",
+    "Tibetan Mastiff",
+    "French Bulldog",
+    "Great Dane",
+    "St. Bernard",
+    "husky",
+    "Alaskan Malamute",
+    "Siberian Husky",
+    "Dalmatian",
+    "Affenpinscher",
+    "Basenji",
+    "pug",
+    "Leonberger",
+    "Newfoundland dog",
+    "Great Pyrenees dog",
+    "Samoyed",
+    "Pomeranian",
+    "Chow Chow",
+    "Keeshond",
+    "brussels griffon",
+    "Pembroke Welsh Corgi",
+    "Cardigan Welsh Corgi",
+    "Toy Poodle",
+    "Miniature Poodle",
+    "Standard Poodle",
+    "Mexican hairless dog (xoloitzcuintli)",
+    "grey wolf",
+    "Alaskan tundra wolf",
+    "red wolf or maned wolf",
+    "coyote",
+    "dingo",
+    "dhole",
+    "African wild dog",
+    "hyena",
+    "red fox",
+    "kit fox",
+    "Arctic fox",
+    "grey fox",
+    "tabby cat",
+    "tiger cat",
+    "Persian cat",
+    "Siamese cat",
+    "Egyptian Mau",
+    "cougar",
+    "lynx",
+    "leopard",
+    "snow leopard",
+    "jaguar",
+    "lion",
+    "tiger",
+    "cheetah",
+    "brown bear",
+    "American black bear",
+    "polar bear",
+    "sloth bear",
+    "mongoose",
+    "meerkat",
+    "tiger beetle",
+    "ladybug",
+    "ground beetle",
+    "longhorn beetle",
+    "leaf beetle",
+    "dung beetle",
+    "rhinoceros beetle",
+    "weevil",
+    "fly",
+    "bee",
+    "ant",
+    "grasshopper",
+    "cricket insect",
+    "stick insect",
+    "cockroach",
+    "praying mantis",
+    "cicada",
+    "leafhopper",
+    "lacewing",
+    "dragonfly",
+    "damselfly",
+    "red admiral butterfly",
+    "ringlet butterfly",
+    "monarch butterfly",
+    "small white butterfly",
+    "sulphur butterfly",
+    "gossamer-winged butterfly",
+    "starfish",
+    "sea urchin",
+    "sea cucumber",
+    "cottontail rabbit",
+    "hare",
+    "Angora rabbit",
+    "hamster",
+    "porcupine",
+    "fox squirrel",
+    "marmot",
+    "beaver",
+    "guinea pig",
+    "common sorrel horse",
+    "zebra",
+    "pig",
+    "wild boar",
+    "warthog",
+    "hippopotamus",
+    "ox",
+    "water buffalo",
+    "bison",
+    "ram (adult male sheep)",
+    "bighorn sheep",
+    "Alpine ibex",
+    "hartebeest",
+    "impala (antelope)",
+    "gazelle",
+    "arabian camel",
+    "llama",
+    "weasel",
+    "mink",
+    "European polecat",
+    "black-footed ferret",
+    "otter",
+    "skunk",
+    "badger",
+    "armadillo",
+    "three-toed sloth",
+    "orangutan",
+    "gorilla",
+    "chimpanzee",
+    "gibbon",
+    "siamang",
+    "guenon",
+    "patas monkey",
+    "baboon",
+    "macaque",
+    "langur",
+    "black-and-white colobus",
+    "proboscis monkey",
+    "marmoset",
+    "white-headed capuchin",
+    "howler monkey",
+    "titi monkey",
+    "Geoffroy's spider monkey",
+    "common squirrel monkey",
+    "ring-tailed lemur",
+    "indri",
+    "Asian elephant",
+    "African bush elephant",
+    "red panda",
+    "giant panda",
+    "snoek fish",
+    "eel",
+    "silver salmon",
+    "rock beauty fish",
+    "clownfish",
+    "sturgeon",
+    "gar fish",
+    "lionfish",
+    "pufferfish",
+    "abacus",
+    "abaya",
+    "academic gown",
+    "accordion",
+    "acoustic guitar",
+    "aircraft carrier",
+    "airliner",
+    "airship",
+    "altar",
+    "ambulance",
+    "amphibious vehicle",
+    "analog clock",
+    "apiary",
+    "apron",
+    "trash can",
+    "assault rifle",
+    "backpack",
+    "bakery",
+    "balance beam",
+    "balloon",
+    "ballpoint pen",
+    "Band-Aid",
+    "banjo",
+    "baluster / handrail",
+    "barbell",
+    "barber chair",
+    "barbershop",
+    "barn",
+    "barometer",
+    "barrel",
+    "wheelbarrow",
+    "baseball",
+    "basketball",
+    "bassinet",
+    "bassoon",
+    "swimming cap",
+    "bath towel",
+    "bathtub",
+    "station wagon",
+    "lighthouse",
+    "beaker",
+    "military hat (bearskin or shako)",
+    "beer bottle",
+    "beer glass",
+    "bell tower",
+    "baby bib",
+    "tandem bicycle",
+    "bikini",
+    "ring binder",
+    "binoculars",
+    "birdhouse",
+    "boathouse",
+    "bobsleigh",
+    "bolo tie",
+    "poke bonnet",
+    "bookcase",
+    "bookstore",
+    "bottle cap",
+    "hunting bow",
+    "bow tie",
+    "brass memorial plaque",
+    "bra",
+    "breakwater",
+    "breastplate",
+    "broom",
+    "bucket",
+    "buckle",
+    "bulletproof vest",
+    "high-speed train",
+    "butcher shop",
+    "taxicab",
+    "cauldron",
+    "candle",
+    "cannon",
+    "canoe",
+    "can opener",
+    "cardigan",
+    "car mirror",
+    "carousel",
+    "tool kit",
+    "cardboard box / carton",
+    "car wheel",
+    "automated teller machine",
+    "cassette",
+    "cassette player",
+    "castle",
+    "catamaran",
+    "CD player",
+    "cello",
+    "mobile phone",
+    "chain",
+    "chain-link fence",
+    "chain mail",
+    "chainsaw",
+    "storage chest",
+    "chiffonier",
+    "bell or wind chime",
+    "china cabinet",
+    "Christmas stocking",
+    "church",
+    "movie theater",
+    "cleaver",
+    "cliff dwelling",
+    "cloak",
+    "clogs",
+    "cocktail shaker",
+    "coffee mug",
+    "coffeemaker",
+    "spiral or coil",
+    "combination lock",
+    "computer keyboard",
+    "candy store",
+    "container ship",
+    "convertible",
+    "corkscrew",
+    "cornet",
+    "cowboy boot",
+    "cowboy hat",
+    "cradle",
+    "construction crane",
+    "crash helmet",
+    "crate",
+    "infant bed",
+    "Crock Pot",
+    "croquet ball",
+    "crutch",
+    "cuirass",
+    "dam",
+    "desk",
+    "desktop computer",
+    "rotary dial telephone",
+    "diaper",
+    "digital clock",
+    "digital watch",
+    "dining table",
+    "dishcloth",
+    "dishwasher",
+    "disc brake",
+    "dock",
+    "dog sled",
+    "dome",
+    "doormat",
+    "drilling rig",
+    "drum",
+    "drumstick",
+    "dumbbell",
+    "Dutch oven",
+    "electric fan",
+    "electric guitar",
+    "electric locomotive",
+    "entertainment center",
+    "envelope",
+    "espresso machine",
+    "face powder",
+    "feather boa",
+    "filing cabinet",
+    "fireboat",
+    "fire truck",
+    "fire screen",
+    "flagpole",
+    "flute",
+    "folding chair",
+    "football helmet",
+    "forklift",
+    "fountain",
+    "fountain pen",
+    "four-poster bed",
+    "freight car",
+    "French horn",
+    "frying pan",
+    "fur coat",
+    "garbage truck",
+    "gas mask or respirator",
+    "gas pump",
+    "goblet",
+    "go-kart",
+    "golf ball",
+    "golf cart",
+    "gondola",
+    "gong",
+    "gown",
+    "grand piano",
+    "greenhouse",
+    "radiator grille",
+    "grocery store",
+    "guillotine",
+    "hair clip",
+    "hair spray",
+    "half-track",
+    "hammer",
+    "hamper",
+    "hair dryer",
+    "hand-held computer",
+    "handkerchief",
+    "hard disk drive",
+    "harmonica",
+    "harp",
+    "combine harvester",
+    "hatchet",
+    "holster",
+    "home theater",
+    "honeycomb",
+    "hook",
+    "hoop skirt",
+    "gymnastic horizontal bar",
+    "horse-drawn vehicle",
+    "hourglass",
+    "iPod",
+    "clothes iron",
+    "carved pumpkin",
+    "jeans",
+    "jeep",
+    "T-shirt",
+    "jigsaw puzzle",
+    "rickshaw",
+    "joystick",
+    "kimono",
+    "knee pad",
+    "knot",
+    "lab coat",
+    "ladle",
+    "lampshade",
+    "laptop computer",
+    "lawn mower",
+    "lens cap",
+    "letter opener",
+    "library",
+    "lifeboat",
+    "lighter",
+    "limousine",
+    "ocean liner",
+    "lipstick",
+    "slip-on shoe",
+    "lotion",
+    "music speaker",
+    "loupe magnifying glass",
+    "sawmill",
+    "magnetic compass",
+    "messenger bag",
+    "mailbox",
+    "tights",
+    "one-piece bathing suit",
+    "manhole cover",
+    "maraca",
+    "marimba",
+    "mask",
+    "matchstick",
+    "maypole",
+    "maze",
+    "measuring cup",
+    "medicine cabinet",
+    "megalith",
+    "microphone",
+    "microwave oven",
+    "military uniform",
+    "milk can",
+    "minibus",
+    "miniskirt",
+    "minivan",
+    "missile",
+    "mitten",
+    "mixing bowl",
+    "mobile home",
+    "ford model t",
+    "modem",
+    "monastery",
+    "monitor",
+    "moped",
+    "mortar and pestle",
+    "graduation cap",
+    "mosque",
+    "mosquito net",
+    "vespa",
+    "mountain bike",
+    "tent",
+    "computer mouse",
+    "mousetrap",
+    "moving van",
+    "muzzle",
+    "metal nail",
+    "neck brace",
+    "necklace",
+    "baby pacifier",
+    "notebook computer",
+    "obelisk",
+    "oboe",
+    "ocarina",
+    "odometer",
+    "oil filter",
+    "pipe organ",
+    "oscilloscope",
+    "overskirt",
+    "bullock cart",
+    "oxygen mask",
+    "product packet / packaging",
+    "paddle",
+    "paddle wheel",
+    "padlock",
+    "paintbrush",
+    "pajamas",
+    "palace",
+    "pan flute",
+    "paper towel",
+    "parachute",
+    "parallel bars",
+    "park bench",
+    "parking meter",
+    "railroad car",
+    "patio",
+    "payphone",
+    "pedestal",
+    "pencil case",
+    "pencil sharpener",
+    "perfume",
+    "Petri dish",
+    "photocopier",
+    "plectrum",
+    "Pickelhaube",
+    "picket fence",
+    "pickup truck",
+    "pier",
+    "piggy bank",
+    "pill bottle",
+    "pillow",
+    "ping-pong ball",
+    "pinwheel",
+    "pirate ship",
+    "drink pitcher",
+    "block plane",
+    "planetarium",
+    "plastic bag",
+    "plate rack",
+    "farm plow",
+    "plunger",
+    "Polaroid camera",
+    "pole",
+    "police van",
+    "poncho",
+    "pool table",
+    "soda bottle",
+    "plant pot",
+    "potter's wheel",
+    "power drill",
+    "prayer rug",
+    "printer",
+    "prison",
+    "missile",
+    "projector",
+    "hockey puck",
+    "punching bag",
+    "purse",
+    "quill",
+    "quilt",
+    "race car",
+    "racket",
+    "radiator",
+    "radio",
+    "radio telescope",
+    "rain barrel",
+    "recreational vehicle",
+    "fishing casting reel",
+    "reflex camera",
+    "refrigerator",
+    "remote control",
+    "restaurant",
+    "revolver",
+    "rifle",
+    "rocking chair",
+    "rotisserie",
+    "eraser",
+    "rugby ball",
+    "ruler measuring stick",
+    "sneaker",
+    "safe",
+    "safety pin",
+    "salt shaker",
+    "sandal",
+    "sarong",
+    "saxophone",
+    "scabbard",
+    "weighing scale",
+    "school bus",
+    "schooner",
+    "scoreboard",
+    "CRT monitor",
+    "screw",
+    "screwdriver",
+    "seat belt",
+    "sewing machine",
+    "shield",
+    "shoe store",
+    "shoji screen / room divider",
+    "shopping basket",
+    "shopping cart",
+    "shovel",
+    "shower cap",
+    "shower curtain",
+    "ski",
+    "balaclava ski mask",
+    "sleeping bag",
+    "slide rule",
+    "sliding door",
+    "slot machine",
+    "snorkel",
+    "snowmobile",
+    "snowplow",
+    "soap dispenser",
+    "soccer ball",
+    "sock",
+    "solar thermal collector",
+    "sombrero",
+    "soup bowl",
+    "keyboard space bar",
+    "space heater",
+    "space shuttle",
+    "spatula",
+    "motorboat",
+    "spider web",
+    "spindle",
+    "sports car",
+    "spotlight",
+    "stage",
+    "steam locomotive",
+    "through arch bridge",
+    "steel drum",
+    "stethoscope",
+    "scarf",
+    "stone wall",
+    "stopwatch",
+    "stove",
+    "strainer",
+    "tram",
+    "stretcher",
+    "couch",
+    "stupa",
+    "submarine",
+    "suit",
+    "sundial",
+    "sunglasses",
+    "sunglasses",
+    "sunscreen",
+    "suspension bridge",
+    "mop",
+    "sweatshirt",
+    "swim trunks / shorts",
+    "swing",
+    "electrical switch",
+    "syringe",
+    "table lamp",
+    "tank",
+    "tape player",
+    "teapot",
+    "teddy bear",
+    "television",
+    "tennis ball",
+    "thatched roof",
+    "front curtain",
+    "thimble",
+    "threshing machine",
+    "throne",
+    "tile roof",
+    "toaster",
+    "tobacco shop",
+    "toilet seat",
+    "torch",
+    "totem pole",
+    "tow truck",
+    "toy store",
+    "tractor",
+    "semi-trailer truck",
+    "tray",
+    "trench coat",
+    "tricycle",
+    "trimaran",
+    "tripod",
+    "triumphal arch",
+    "trolleybus",
+    "trombone",
+    "hot tub",
+    "turnstile",
+    "typewriter keyboard",
+    "umbrella",
+    "unicycle",
+    "upright piano",
+    "vacuum cleaner",
+    "vase",
+    "vaulted or arched ceiling",
+    "velvet fabric",
+    "vending machine",
+    "vestment",
+    "viaduct",
+    "violin",
+    "volleyball",
+    "waffle iron",
+    "wall clock",
+    "wallet",
+    "wardrobe",
+    "military aircraft",
+    "sink",
+    "washing machine",
+    "water bottle",
+    "water jug",
+    "water tower",
+    "whiskey jug",
+    "whistle",
+    "hair wig",
+    "window screen",
+    "window shade",
+    "Windsor tie",
+    "wine bottle",
+    "airplane wing",
+    "wok",
+    "wooden spoon",
+    "wool",
+    "split-rail fence",
+    "shipwreck",
+    "sailboat",
+    "yurt",
+    "website",
+    "comic book",
+    "crossword",
+    "traffic or street sign",
+    "traffic light",
+    "dust jacket",
+    "menu",
+    "plate",
+    "guacamole",
+    "consomme",
+    "hot pot",
+    "trifle",
+    "ice cream",
+    "popsicle",
+    "baguette",
+    "bagel",
+    "pretzel",
+    "cheeseburger",
+    "hot dog",
+    "mashed potatoes",
+    "cabbage",
+    "broccoli",
+    "cauliflower",
+    "zucchini",
+    "spaghetti squash",
+    "acorn squash",
+    "butternut squash",
+    "cucumber",
+    "artichoke",
+    "bell pepper",
+    "cardoon",
+    "mushroom",
+    "Granny Smith apple",
+    "strawberry",
+    "orange",
+    "lemon",
+    "fig",
+    "pineapple",
+    "banana",
+    "jackfruit",
+    "cherimoya (custard apple)",
+    "pomegranate",
+    "hay",
+    "carbonara",
+    "chocolate syrup",
+    "dough",
+    "meatloaf",
+    "pizza",
+    "pot pie",
+    "burrito",
+    "red wine",
+    "espresso",
+    "tea cup",
+    "eggnog",
+    "mountain",
+    "bubble",
+    "cliff",
+    "coral reef",
+    "geyser",
+    "lakeshore",
+    "promontory",
+    "sandbar",
+    "beach",
+    "valley",
+    "volcano",
+    "baseball player",
+    "bridegroom",
+    "scuba diver",
+    "rapeseed",
+    "daisy",
+    "yellow lady's slipper",
+    "corn",
+    "acorn",
+    "rose hip",
+    "horse chestnut seed",
+    "coral fungus",
+    "agaric",
+    "gyromitra",
+    "stinkhorn mushroom",
+    "earth star fungus",
+    "hen of the woods mushroom",
+    "bolete",
+    "corn cob",
+    "toilet paper",
+]
+
+
+openai_imagenet_template = [
+    lambda c: f"a bad photo of a {c}.",
+    lambda c: f"a photo of many {c}.",
+    lambda c: f"a sculpture of a {c}.",
+    lambda c: f"a photo of the hard to see {c}.",
+    lambda c: f"a low resolution photo of the {c}.",
+    lambda c: f"a rendering of a {c}.",
+    lambda c: f"graffiti of a {c}.",
+    lambda c: f"a bad photo of the {c}.",
+    lambda c: f"a cropped photo of the {c}.",
+    lambda c: f"a tattoo of a {c}.",
+    lambda c: f"the embroidered {c}.",
+    lambda c: f"a photo of a hard to see {c}.",
+    lambda c: f"a bright photo of a {c}.",
+    lambda c: f"a photo of a clean {c}.",
+    lambda c: f"a photo of a dirty {c}.",
+    lambda c: f"a dark photo of the {c}.",
+    lambda c: f"a drawing of a {c}.",
+    lambda c: f"a photo of my {c}.",
+    lambda c: f"the plastic {c}.",
+    lambda c: f"a photo of the cool {c}.",
+    lambda c: f"a close-up photo of a {c}.",
+    lambda c: f"a black and white photo of the {c}.",
+    lambda c: f"a painting of the {c}.",
+    lambda c: f"a painting of a {c}.",
+    lambda c: f"a pixelated photo of the {c}.",
+    lambda c: f"a sculpture of the {c}.",
+    lambda c: f"a bright photo of the {c}.",
+    lambda c: f"a cropped photo of a {c}.",
+    lambda c: f"a plastic {c}.",
+    lambda c: f"a photo of the dirty {c}.",
+    lambda c: f"a jpeg corrupted photo of a {c}.",
+    lambda c: f"a blurry photo of the {c}.",
+    lambda c: f"a photo of the {c}.",
+    lambda c: f"a good photo of the {c}.",
+    lambda c: f"a rendering of the {c}.",
+    lambda c: f"a {c} in a video game.",
+    lambda c: f"a photo of one {c}.",
+    lambda c: f"a doodle of a {c}.",
+    lambda c: f"a close-up photo of the {c}.",
+    lambda c: f"a photo of a {c}.",
+    lambda c: f"the origami {c}.",
+    lambda c: f"the {c} in a video game.",
+    lambda c: f"a sketch of a {c}.",
+    lambda c: f"a doodle of the {c}.",
+    lambda c: f"a origami {c}.",
+    lambda c: f"a low resolution photo of a {c}.",
+    lambda c: f"the toy {c}.",
+    lambda c: f"a rendition of the {c}.",
+    lambda c: f"a photo of the clean {c}.",
+    lambda c: f"a photo of a large {c}.",
+    lambda c: f"a rendition of a {c}.",
+    lambda c: f"a photo of a nice {c}.",
+    lambda c: f"a photo of a weird {c}.",
+    lambda c: f"a blurry photo of a {c}.",
+    lambda c: f"a cartoon {c}.",
+    lambda c: f"art of a {c}.",
+    lambda c: f"a sketch of the {c}.",
+    lambda c: f"a embroidered {c}.",
+    lambda c: f"a pixelated photo of a {c}.",
+    lambda c: f"itap of the {c}.",
+    lambda c: f"a jpeg corrupted photo of the {c}.",
+    lambda c: f"a good photo of a {c}.",
+    lambda c: f"a plushie {c}.",
+    lambda c: f"a photo of the nice {c}.",
+    lambda c: f"a photo of the small {c}.",
+    lambda c: f"a photo of the weird {c}.",
+    lambda c: f"the cartoon {c}.",
+    lambda c: f"art of the {c}.",
+    lambda c: f"a drawing of the {c}.",
+    lambda c: f"a photo of the large {c}.",
+    lambda c: f"a black and white photo of a {c}.",
+    lambda c: f"the plushie {c}.",
+    lambda c: f"a dark photo of a {c}.",
+    lambda c: f"itap of a {c}.",
+    lambda c: f"graffiti of the {c}.",
+    lambda c: f"a toy {c}.",
+    lambda c: f"itap of my {c}.",
+    lambda c: f"a photo of a cool {c}.",
+    lambda c: f"a photo of a small {c}.",
+    lambda c: f"a tattoo of the {c}.",
+]

multimodal/examples/flava/data/transforms.py

@@ -0,0 +1,131 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+from functools import partial
+from typing import Any, Callable, Optional
+
+import torch
+from torchmultimodal.transforms.flava_transform import FLAVAImageTransform
+from torchvision import transforms
+from transformers import BertTokenizer
+
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+IMAGE_DEFAULT_SIZE = (224, 224)
+VL_MAX_LENGTH_DEFAULT = 77
+TEXT_MAX_LENGTH_DEFAULT = 512
+TEXT_DEFAULT_TOKENIZER = "bert-base-uncased"
+TEXT_WHOLE_WORD_MASK_TOKENIZER = "bert-large-uncased-whole-word-masking"
+
+
+def encode_text(text, tokenizer, *args, **kwargs):
+    return tokenizer(text, *args, **kwargs)
+
+
+def encode_text_batch(
+    batch, tokenizer, text_columns, return_batch=False, *args, **kwargs
+):
+    texts = [batch[column] for column in text_columns]
+    tokens = tokenizer(*texts, *args, **kwargs)
+    if return_batch:
+        batch.update(tokens)
+        return batch
+    return tokens
+
+
+def transform_image_dict(transform, image_dict, *args, **kwargs):
+    return {"image": transform(image_dict["image"], *args, **kwargs)}
+
+
+def default_torchvision_transforms(
+    size=IMAGE_DEFAULT_SIZE,
+    mean=IMAGENET_DEFAULT_MEAN,
+    std=IMAGENET_DEFAULT_STD,
+    use_dict=False,
+):
+    transform = transforms.Compose(
+        [
+            transforms.Resize(size),
+            transforms.ToTensor(),
+            transforms.Normalize(
+                mean=mean,
+                std=std,
+            ),
+        ]
+    )
+
+    if use_dict:
+        transform = partial(transform_image_dict, transform=transform)
+
+    return transform, transform
+
+
+def default_image_pretraining_transforms():
+    return FLAVAImageTransform(), FLAVAImageTransform(is_train=False)
+
+
+def default_text_transform(
+    text_tokenizer: Optional[Callable] = None,
+    max_text_length: int = TEXT_MAX_LENGTH_DEFAULT,
+    **kwargs: Any,
+):
+    if text_tokenizer is None:
+        text_tokenizer = BertTokenizer.from_pretrained(TEXT_DEFAULT_TOKENIZER)
+
+    text_transform = partial(
+        encode_text,
+        tokenizer=text_tokenizer,
+        padding="max_length",
+        max_length=max_text_length,
+        truncation=True,
+        return_tensors="pt",
+        return_special_tokens_mask=True,
+    )
+
+    return text_transform
+
+
+def default_vl_text_transform(
+    text_tokenizer: Optional[Callable] = None,
+    max_text_length: int = VL_MAX_LENGTH_DEFAULT,
+    **kwargs: Any,
+):
+    if text_tokenizer is None:
+        text_tokenizer = BertTokenizer.from_pretrained(TEXT_WHOLE_WORD_MASK_TOKENIZER)
+    return default_text_transform(text_tokenizer, max_text_length=max_text_length)
+
+
+def pad_batch(batch, batch_size):
+    for item in batch.keys():
+        if isinstance(batch[item], torch.Tensor):
+            diff = batch_size - batch[item].size(0)
+            pad = batch[item][-diff:].detach().clone()
+            batch[item] = torch.cat([batch[item], pad], dim=0)
+    return batch
+
+
+class VLTransform:
+    def __init__(self, image_transform, text_transform):
+        self.image_transform = image_transform
+        self.text_transform = text_transform
+
+    def __call__(self, info, dataset, itm_probability):
+        output = {}
+        text = info["text"]
+        image = info["image"]
+        if itm_probability > 0:
+            output["itm_labels"] = torch.ones((1), dtype=torch.long)
+
+        if random.random() < itm_probability:
+            while text == info["text"]:
+                text = dataset.select([random.randint(0, len(dataset) - 1)])[0]["text"]
+            output["itm_labels"] = torch.zeros((1), dtype=torch.long)
+
+        output.update(self.image_transform(image))
+        output.update(self.text_transform(text))
+        return output

multimodal/examples/flava/data/utils.py

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import time
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from typing import List
+
+import requests
+from datasets import concatenate_datasets, load_dataset
+from datasets.utils.file_utils import get_datasets_user_agent
+from flava.definitions import HFDatasetInfo
+from PIL import Image, UnidentifiedImageError
+
+
+DATASETS_USER_AGENT = get_datasets_user_agent()
+
+
+def build_datasets_from_info(dataset_infos: List[HFDatasetInfo], split: str = "train"):
+    dataset_list = []
+    for dataset_info in dataset_infos:
+        current_dataset = load_dataset(
+            dataset_info.key,
+            dataset_info.subset,
+            split=dataset_info.split_key_mapping[split],
+            use_auth_token=True,
+            **dataset_info.extra_kwargs,
+        )
+        if dataset_info.remove_columns is not None:
+            current_dataset = current_dataset.remove_columns(
+                dataset_info.remove_columns
+            )
+        if dataset_info.rename_columns is not None:
+            for rename in dataset_info.rename_columns:
+                current_dataset = current_dataset.rename_column(rename[0], rename[1])
+
+        dataset_list.append(current_dataset)
+
+    return concatenate_datasets(dataset_list)
+
+
+def fetch_single_image(image_url, timeout, retries=0, sleep_timer=0):
+    for _ in range(retries + 1):
+        try:
+            image = Image.open(
+                requests.get(
+                    image_url,
+                    stream=True,
+                    headers={"user-agent": DATASETS_USER_AGENT},
+                    timeout=timeout,
+                ).raw
+            )
+            break
+        except (requests.exceptions.ConnectionError, UnidentifiedImageError):
+            image = None
+            time.sleep(sleep_timer)
+
+    return image
+
+
+def fetch_images(batch, num_threads, timeout=None, retries=0, sleep_timer=0):
+    if "image" in batch:
+        # This dataset already has "image" defined.
+        return batch
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        batch["image"] = list(
+            executor.map(
+                partial(
+                    fetch_single_image,
+                    timeout=timeout,
+                    retries=retries,
+                    sleep_timer=sleep_timer,
+                ),
+                batch["image_url"],
+            )
+        )
+    return batch

multimodal/examples/flava/native/README.md

@@ -0,0 +1,43 @@
+# Usage Instructions
+
+This is a lightweight native pytorch implementation to run scaling studies on the FLAVA model. The original code is located at: [`examples/flava/train.py`](https://github.com/facebookresearch/multimodal/blob/main/examples/flava/train.py)
+
+## Prerequisites
+
+- Install torchmultimodal library [from source](https://github.com/facebookresearch/multimodal/blob/main/README.md#building-from-source)
+- `cd multimodal/examples`
+- `pip install -r flava/requirements.txt`
+
+## Training
+
+### Configuration
+
+Configuration presets for various model sizes can be found at: `examples/flava/native/configs`
+
+Some config settings that are relevant for scaling: (local) `batch_size`, `activation_checkpointing`, `strategy`.
+
+Configs can be overridden through command line, for example: `python -m flava.native.train config=flava/native/configs/pretrain_debug.yaml training.batch_size=8 training.enable_amp=True training.activation_checkpointing=True training.strategy=fsdp`
+
+### Running
+
+
+Using [`torchrun`](https://pytorch.org/docs/stable/elastic/run.html):
+
+**Single node**
+
+`NUM_GPUS=8; torchrun --nproc_per_node=$NUM_GPUS -m flava.native.train config=flava/native/configs/pretrain_debug.yaml`
+
+**Multiple nodes (using slurm)**
+
+Create a `run.slurm` file:
+
+```bash
+RDZV_ENDPOINT=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
+
+srun torchrun --nnodes=$SLURM_NNODES --nproc_per_node=$SLURM_GPUS_PER_TASK --rdzv_id=$SLURM_JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$RDZV_ENDPOINT --max_restarts 0  -m flava.native.train config=flava/native/configs/pretrain_debug.yaml
+$@
+```
+
+Run in terminal:
+
+`sbatch --partition=[PARTITION] --nodes=[NUM_NODES] --gpus-per-task=[NUM_GPUS_PER_NODE] run.slurm`

multimodal/examples/flava/native/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

multimodal/examples/flava/native/configs/1.8b.yaml

@@ -0,0 +1,79 @@
+training:
+  strategy: fsdp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 8
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: True
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train
+
+model:
+  image_num_hidden_layers: 32
+  image_hidden_size: 1280
+  image_intermediate_size: 5120
+  image_num_attention_heads: 16
+
+  text_num_hidden_layers: 32
+  text_hidden_size: 1280
+  text_intermediate_size: 5120
+  text_num_attention_heads: 16
+
+  multimodal_num_hidden_layers: 16
+  multimodal_hidden_size: 1280
+  multimodal_intermediate_size: 5120
+  multimodal_num_attention_heads: 16

multimodal/examples/flava/native/configs/10b.yaml

@@ -0,0 +1,80 @@
+training:
+  strategy: fsdp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 8
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: True
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train
+
+
+model:
+  image_num_hidden_layers: 64
+  image_hidden_size: 2048
+  image_intermediate_size: 10240
+  image_num_attention_heads: 16
+
+  text_num_hidden_layers: 64
+  text_hidden_size: 2048
+  text_intermediate_size: 10240
+  text_num_attention_heads: 16
+
+  multimodal_num_hidden_layers: 40
+  multimodal_hidden_size: 2048
+  multimodal_intermediate_size: 10240
+  multimodal_num_attention_heads: 16

multimodal/examples/flava/native/configs/2.7b.yaml

@@ -0,0 +1,79 @@
+training:
+  strategy: fsdp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 8
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: True
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train
+
+model:
+  image_num_hidden_layers: 40
+  image_hidden_size: 1408
+  image_intermediate_size: 6144
+  image_num_attention_heads: 16
+
+  text_num_hidden_layers: 40
+  text_hidden_size: 1408
+  text_intermediate_size: 6144
+  text_num_attention_heads: 16
+
+  multimodal_num_hidden_layers: 20
+  multimodal_hidden_size: 1408
+  multimodal_intermediate_size: 6144
+  multimodal_num_attention_heads: 16

multimodal/examples/flava/native/configs/4.8b.yaml

@@ -0,0 +1,79 @@
+training:
+  strategy: fsdp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 12
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: True
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train
+
+model:
+  image_num_hidden_layers: 48
+  image_hidden_size: 1664
+  image_intermediate_size: 8192
+  image_num_attention_heads: 16
+
+  text_num_hidden_layers: 48
+  text_hidden_size: 1664
+  text_intermediate_size: 8192
+  text_num_attention_heads: 16
+
+  multimodal_num_hidden_layers: 24
+  multimodal_hidden_size: 1664
+  multimodal_intermediate_size: 8192
+  multimodal_num_attention_heads: 16

multimodal/examples/flava/native/configs/900m.yaml

@@ -0,0 +1,79 @@
+training:
+  strategy: ddp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 8
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: True
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train
+
+model:
+  image_num_hidden_layers: 24
+  image_hidden_size: 1024
+  image_intermediate_size: 4096
+  image_num_attention_heads: 16
+
+  text_num_hidden_layers: 24
+  text_hidden_size: 1024
+  text_intermediate_size: 4096
+  text_num_attention_heads: 16
+
+  multimodal_num_hidden_layers: 12
+  multimodal_hidden_size: 1024
+  multimodal_intermediate_size: 4096
+  multimodal_num_attention_heads: 16

multimodal/examples/flava/native/configs/pretrain_debug.yaml

@@ -0,0 +1,63 @@
+training:
+  strategy: ddp # can be changed to ddp or fsdp
+  seed: 1337
+
+  batch_size: 8
+  num_workers: 4
+  prefetch_factor: 3
+
+  optimizer:
+    learning_rate: 1e-3
+    adam_eps: 1e-8
+    adam_weight_decay: 0.1
+    adam_betas: [0.9, 0.999]
+
+  warmup_steps: 10000
+  max_steps: 100000
+
+  validation_steps: 5000
+  log_interval: 10
+
+  enable_tf32: True
+  enable_amp: True
+  half_precision_format: "bfloat16"  # or float16
+  enable_half_reduce_in_fsdp: True  # handles the reduction across devices in half precision
+
+  activation_checkpointing: False
+
+datasets:
+  _target_: flava.definitions.TrainingDatasetsInfo
+  selected:
+  - image
+  - vl
+  - text
+  image:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: imagenet-1k
+        subset: default
+  text:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: wikitext
+        subset: wikitext-103-raw-v1
+    datamodule_extra_kwargs:
+      text_columns: ["text"]
+  vl:
+    _target_: flava.definitions.TrainingSingleDatasetInfo
+    train:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+    val:
+      - _target_: flava.definitions.HFDatasetInfo
+        key: red_caps
+        subset: backpacking
+        rename_columns:
+          - ["caption", "text"]
+        split_key_mapping:
+          validation: train

multimodal/examples/flava/native/data.py

@@ -0,0 +1,560 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torchvision
+
+from flava.data.transforms import (
+    default_image_pretraining_transforms,
+    default_text_transform,
+    default_torchvision_transforms,
+    encode_text_batch,
+    pad_batch,
+    TEXT_DEFAULT_TOKENIZER,
+    TEXT_WHOLE_WORD_MASK_TOKENIZER,
+    VL_MAX_LENGTH_DEFAULT,
+    VLTransform,
+)
+from flava.data.utils import build_datasets_from_info, fetch_images
+from flava.definitions import HFDatasetInfo, TorchVisionDatasetInfo
+from pytorch_lightning import LightningDataModule
+from torch.utils.data.distributed import DistributedSampler
+from transformers import (
+    BertTokenizer,
+    DataCollatorForLanguageModeling,
+    DataCollatorForWholeWordMask,
+    DefaultDataCollator,
+    TRANSFORMERS_CACHE,
+)
+from transformers.data.data_collator import torch_default_data_collator
+
+
+def transform_image(transform, sample):
+    sample.update(transform(sample["image"]))
+    return sample
+
+
+def get_sampler(dataset, shuffle=True):
+    if dist.is_initialized():
+        return DistributedSampler(dataset, shuffle=shuffle)
+    if shuffle:
+        return torch.utils.data.RandomSampler(dataset)
+    return torch.utils.data.SequentialSampler(dataset)
+
+
+class DataCollatorForWholeWordMaskRetainingBatch(DataCollatorForWholeWordMask):
+    def torch_call(
+        self, examples: List[Union[List[int], Any, Dict[str, Any]]]
+    ) -> Dict[str, Any]:
+        masked_batch = super().torch_call(examples)
+        examples = torch_default_data_collator(examples)
+        examples["input_ids"] = masked_batch["input_ids"]
+        examples["labels"] = masked_batch["labels"]
+        return examples
+
+
+class ImageDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        transforms: Optional[Tuple[Callable, Callable]] = None,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        allow_uneven_batches: bool = False,
+        prefetch_factor: int = 2,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_dataset_infos = train_infos
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.allow_uneven_batches = allow_uneven_batches
+        self.prefetch_factor = prefetch_factor
+
+        if transforms is None:
+            transforms = default_image_pretraining_transforms()
+
+        self.train_transform, self.test_transform = transforms
+
+    def setup(self, stage=None):
+        train_transform = partial(transform_image, self.train_transform)
+        val_transform = partial(transform_image, self.test_transform)
+
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(train_transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(val_transform)
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=get_sampler(self.train_dataset, shuffle=True),
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            # ideally, we don't need to drop these for unimodal cases
+            # but just to be safe
+            drop_last=True,
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=get_sampler(self.val_dataset, shuffle=False),
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            # ideally, we don't need to drop these for unimodal cases
+            # but just to be safe
+            drop_last=True,
+        )
+
+    def test_dataloader(self):
+        return self.val_dataloader()
+
+    def on_before_batch_transfer(self, batch, *args):
+        if batch["label"].size(0) < self.batch_size and not self.allow_uneven_batches:
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+
+class TextDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        text_columns: List[str],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        tokenizer: Optional[Callable] = None,
+        max_length: int = 512,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        allow_uneven_batches: bool = False,
+        prefetch_factor: int = 2,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_dataset_infos = train_infos
+        self.text_columns = text_columns
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.allow_uneven_batches = allow_uneven_batches
+        self.prefetch_factor = prefetch_factor
+
+    def setup(self, stage=None):
+        if self.tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained(TEXT_DEFAULT_TOKENIZER)
+        transform = partial(
+            encode_text_batch,
+            tokenizer=self.tokenizer,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_tensors="pt",
+            return_special_tokens_mask=True,
+            text_columns=self.text_columns,
+            return_batch=True,
+        )
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(transform)
+
+    def train_dataloader(self):
+        return self._build_dataloader(self.train_dataset)
+
+    def val_dataloader(self):
+        return self._build_dataloader(self.val_dataset, shuffle=False)
+
+    def _build_dataloader(self, dataset, drop_last=False, shuffle=True):
+        return torch.utils.data.DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=get_sampler(dataset, shuffle),
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+            collate_fn=self._build_collator(),
+            drop_last=drop_last,
+        )
+
+    def _build_collator(self):
+        return DefaultDataCollator()
+
+    def on_before_batch_transfer(self, batch, *args):
+        batch.pop("token_type_ids", None)
+        mask = batch.pop("attention_mask", None)
+        if mask.size(0) < self.batch_size and not self.allow_uneven_batches:
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+    def on_after_batch_transfer(self, batch, *args):
+        batch["text"] = batch.pop("input_ids")
+        return batch
+
+
+class MLMDataModule(TextDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        text_columns: List[str],
+        val_infos: Optional[List[HFDatasetInfo]] = None,
+        mlm_probability: float = 0.15,
+        ignore_index: int = -1,
+        **kwargs: Any,
+    ):
+        super().__init__(train_infos, text_columns, val_infos, **kwargs)
+        self.mlm_probability = mlm_probability
+        self.ignore_index = ignore_index
+
+    def setup(self, stage=None):
+        if self.tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained(TEXT_DEFAULT_TOKENIZER)
+        transform = partial(
+            encode_text_batch,
+            tokenizer=self.tokenizer,
+            padding="max_length",
+            max_length=self.max_length,
+            truncation=True,
+            return_tensors="pt",
+            return_special_tokens_mask=True,
+            text_columns=self.text_columns,
+            return_batch=False,
+        )
+        self.train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        self.train_dataset.set_transform(transform)
+        self.val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+        self.val_dataset.set_transform(transform)
+
+    def _build_dataloader(self, dataset, drop_last=True, shuffle=True):
+        # uneven batches can cause distributed issues,
+        # drop last batch to prevent those.
+        # ideally, we don't need to drop these for unimodal cases
+        # but just to be safe
+        return super()._build_dataloader(dataset, drop_last=drop_last, shuffle=shuffle)
+
+    def _build_collator(self):
+        return DataCollatorForLanguageModeling(
+            self.tokenizer, mlm_probability=self.mlm_probability
+        )
+
+    def on_after_batch_transfer(self, batch, *args):
+        batch["text_masked"] = batch.pop("input_ids")
+        batch["mlm_labels"] = batch.pop("labels")
+        batch["mlm_labels"][batch["mlm_labels"] == -100] = self.ignore_index
+        return batch
+
+
+class VLDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[HFDatasetInfo],
+        val_infos: List[HFDatasetInfo],
+        text_transform: Optional[Callable] = None,
+        image_transforms: Optional[Tuple[Callable, Callable]] = None,
+        mlm_probablity: float = 0.15,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        finetuning: bool = False,
+        ignore_index: int = -1,
+        itm_probability: float = 0.1,
+        allow_uneven_batches: bool = False,
+        fetch_num_threads: int = 4,
+        fetch_retries: int = 0,
+        fetch_sleep_timer: int = 0,
+        fetch_timeout: Optional[float] = None,
+        fetch_batch_size: int = 50,
+        prefetch_factor=2,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.train_dataset_infos = train_infos
+        self.val_dataset_infos = val_infos
+        if self.val_dataset_infos is None:
+            self.val_dataset_infos = train_infos
+        if image_transforms is None:
+            if not finetuning:
+                image_transforms = default_image_pretraining_transforms()
+            else:
+                image_transforms = default_torchvision_transforms(use_dict=True)
+
+        self.train_image_transform, self.test_image_transform = image_transforms
+        self.text_transform = text_transform
+        self.mlm_probability = mlm_probablity
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.ignore_index = ignore_index
+        self.itm_probability = itm_probability
+        self.allow_uneven_batches = allow_uneven_batches
+        self.fetch_num_threads = fetch_num_threads
+        self.fetch_retries = fetch_retries
+        self.fetch_sleep_timer = fetch_sleep_timer
+        self.fetch_timeout = fetch_timeout
+        self.fetch_batch_size = fetch_batch_size
+        self.prefetch_factor = prefetch_factor
+
+    def setup(self, stage=None):
+        if self.text_transform is None:
+            # TODO Update to use whole word mask vocab
+            text_tokenizer = BertTokenizer.from_pretrained(
+                TEXT_WHOLE_WORD_MASK_TOKENIZER
+            )
+            self.text_transform = default_text_transform(
+                text_tokenizer, max_text_length=VL_MAX_LENGTH_DEFAULT
+            )
+        self.text_tokenizer = self.text_transform.keywords["tokenizer"]
+        train_vl_transform = VLTransform(
+            self.train_image_transform, self.text_transform
+        )
+        val_vl_transform = VLTransform(self.test_image_transform, self.text_transform)
+
+        train_dataset = build_datasets_from_info(
+            self.train_dataset_infos, split="train"
+        )
+        train_dataset = train_dataset.map(
+            fetch_images,
+            batched=True,
+            batch_size=self.fetch_batch_size,
+            fn_kwargs={
+                "num_threads": self.fetch_num_threads,
+                "timeout": self.fetch_timeout,
+                "retries": self.fetch_retries,
+                "sleep_timer": self.fetch_sleep_timer,
+            },
+        )
+        train_dataset = train_dataset.filter(
+            lambda example: example["image"] is not None
+        )
+        self.train_dataset = train_dataset
+        self.train_dataset.set_transform(
+            partial(
+                train_vl_transform,
+                dataset=train_dataset.filter(lambda example: True),
+                itm_probability=self.itm_probability,
+            )
+        )
+
+        val_dataset = build_datasets_from_info(
+            self.val_dataset_infos, split="validation"
+        )
+
+        val_dataset = val_dataset.map(
+            fetch_images,
+            batched=True,
+            batch_size=self.fetch_batch_size,
+            fn_kwargs={
+                "num_threads": self.fetch_num_threads,
+                "timeout": self.fetch_timeout,
+                "retries": self.fetch_retries,
+                "sleep_timer": self.fetch_sleep_timer,
+            },
+        )
+        val_dataset = val_dataset.filter(lambda example: example["image"] is not None)
+        self.val_dataset = val_dataset
+        self.val_dataset.set_transform(
+            partial(
+                val_vl_transform,
+                dataset=self.val_dataset.filter(
+                    lambda example: True
+                ),  # Pass a copy to transform
+                itm_probability=self.itm_probability,
+            )
+        )
+
+    def train_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=get_sampler(self.train_dataset),
+            collate_fn=self._build_collator(),
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            drop_last=True,
+        )
+
+    def val_dataloader(self):
+        return torch.utils.data.DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            sampler=get_sampler(self.val_dataset, shuffle=False),
+            collate_fn=self._build_collator(),
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+            # uneven batches can cause distributed issues,
+            # drop last batch to prevent those.
+            drop_last=True,
+        )
+
+    def _build_collator(self):
+        return DataCollatorForWholeWordMaskRetainingBatch(
+            self.text_tokenizer, mlm_probability=self.mlm_probability
+        )
+
+    def on_before_batch_transfer(self, batch, *args):
+        batch.pop("token_type_ids", None)
+        mask = batch.pop("attention_mask", None)
+        if (
+            mask is not None
+            and mask.size(0) < self.batch_size
+            and not self.allow_uneven_batches
+        ):
+            batch = pad_batch(batch, self.batch_size)
+        return batch
+
+    def on_after_batch_transfer(self, batch, *args):
+        text_masked = batch.pop("input_ids")
+        mlm_labels = batch.pop("labels", None)
+        mlm_labels[mlm_labels == -100] = self.ignore_index
+        text = text_masked.detach().clone()
+        text[mlm_labels != -1] = mlm_labels[mlm_labels != -1]
+        batch.update(
+            {"mlm_labels": mlm_labels, "text": text, "text_masked": text_masked}
+        )
+        return batch
+
+
+class TorchVisionDataModule(LightningDataModule):
+    def __init__(
+        self,
+        train_infos: List[TorchVisionDatasetInfo],
+        # Val info is not used for torchvision datamodule, but kept to keep things consistent
+        val_infos: Optional[List[TorchVisionDatasetInfo]] = None,
+        dataset_root: Optional[str] = None,
+        image_transforms: Optional[Tuple[Callable, Callable]] = None,
+        batch_size: int = 32,
+        num_workers: int = 4,
+        prefetch_factor: int = 2,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        self.train_info = train_infos[0]
+        if val_infos is None:
+            val_infos = train_infos
+        self.val_info = val_infos[0]
+
+        self.train_class_ptr, self.train_root = self._parse_info(
+            self.train_info, dataset_root=dataset_root
+        )
+        self.val_class_ptr, self.val_root = self._parse_info(
+            self.val_info, dataset_root=dataset_root
+        )
+
+        if image_transforms is None:
+            image_transforms = default_torchvision_transforms()
+
+        self.train_transform, self.test_transform = image_transforms
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.prefetch_factor = prefetch_factor
+
+    def _parse_info(
+        self, info: TorchVisionDatasetInfo, dataset_root: Optional[str] = None
+    ):
+        assert hasattr(
+            torchvision.datasets, info.key
+        ), f"No dataset named {info.key} present in torchvision.datasets"
+        class_ptr = getattr(torchvision.datasets, info.key)
+        if dataset_root is None:
+            dataset_root = os.path.join(TRANSFORMERS_CACHE, "datasets", "torchvision")
+            dataset_root = os.path.join(dataset_root, class_ptr.__name__.lower())
+            os.makedirs(dataset_root, exist_ok=True)
+
+        return class_ptr, dataset_root
+
+    def setup(self, stage=None):
+        self.train_dataset = self.train_class_ptr(
+            self.train_root,
+            split=self.train_info.train_split,
+            transform=self.train_transform,
+            download=True,
+        )
+
+        if self.val_info.has_val:
+            self.val_dataset = self.val_class_ptr(
+                self.val_root,
+                split=self.val_info.val_split,
+                transform=self.test_transform,
+                download=True,
+            )
+
+        self.test_dataset = self.val_class_ptr(
+            self.val_root,
+            split=self.val_info.test_split,
+            transform=self.test_transform,
+            download=True,
+        )
+
+    def train_dataloader(self):
+        return self._build_dataloader(self.train_dataset)
+
+    def val_dataloader(self):
+        if self.val_info.has_val:
+            dataset = self.val_dataset
+        else:
+            dataset = self.test_dataset
+
+        return self._build_dataloader(dataset, shuffle=False)
+
+    def test_dataloader(self):
+        return self._build_dataloader(self.test_dataset, shuffle=False)
+
+    def _build_dataloader(self, dataset: torch.utils.data.Dataset, shuffle=True):
+        return torch.utils.data.DataLoader(
+            dataset,
+            sampler=get_sampler(dataset, shuffle),
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=True,
+            prefetch_factor=self.prefetch_factor,
+        )
+
+    def on_before_batch_transfer(self, batch, *args):
+        images, targets = batch
+        batch = {"image": images, "labels": targets}
+        return batch

multimodal/examples/flava/native/model.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, Tuple
+
+import torch
+from torch import nn
+from torchmultimodal.models.flava.model import flava_model_for_pretraining
+from transformers.optimization import get_cosine_schedule_with_warmup
+
+
+def get_optimizer(
+    model: torch.nn.Module,
+    learning_rate: float = 0.0002,
+    adam_eps: float = 1.0e-08,
+    adam_weight_decay: float = 0.01,
+    adam_betas: Tuple[int, int] = (0.9, 0.999),
+    warmup_steps: int = 2000,
+    max_steps: int = 450000,
+):
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=learning_rate,
+        betas=adam_betas,
+        eps=adam_eps,
+        weight_decay=adam_weight_decay,
+    )
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=warmup_steps,
+        num_training_steps=max_steps,
+    )
+    return optimizer, scheduler
+
+
+class FLAVAPreTrainModule(nn.Module):
+    def __init__(
+        self,
+        use_bf16: bool = True,
+        **flava_pretraining_kwargs: Any,
+    ):
+        super().__init__()
+        self.model = flava_model_for_pretraining(**flava_pretraining_kwargs)
+        self.use_bf16 = use_bf16
+
+    def forward(self, batch, action=None):
+        # super hacky
+        if action == "encode_text":
+            return self.model.encode_text(batch)
+        elif action == "encode_image":
+            return self.model.encode_image(batch)
+
+        if "image" in batch and ("text" in batch or "text_masked" in batch):
+            required_embedding = "mm"
+        elif "image" in batch:
+            required_embedding = "image"
+        elif "text" in batch or "text_masked" in batch:
+            required_embedding = "text"
+        else:
+            raise RuntimeError("Batch needs to have either or both 'image' and 'text'.")
+
+        output = self.model(
+            image=batch.get("image"),
+            image_for_codebook=batch.get("image_for_codebook"),
+            image_patches_mask=batch.get("image_patches_mask"),
+            text=batch.get("text"),
+            text_masked=batch.get("text_masked"),
+            mlm_labels=batch.get("mlm_labels"),
+            itm_labels=batch.get("itm_labels"),
+            required_embedding=required_embedding,
+        )
+        return output
+
+    def encode_text(self, *args, **kwargs):
+        return self.model.encode_text(*args, **kwargs)

multimodal/examples/flava/native/train.py

@@ -0,0 +1,415 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# example command to train:
+# `torchrun --nproc_per_node=8 -m flava.native.train config=flava/native/configs/pretrain_debug.yaml`
+
+
+import time
+from functools import partial
+from typing import Any, Dict, Tuple, Union
+
+import datasets
+import numpy as np
+import torch
+import torch.distributed as dist
+from common.data import MultiDataModule
+from flava.definitions import FLAVAArguments
+from flava.native.data import (
+    default_text_transform,
+    ImageDataModule,
+    MLMDataModule,
+    VL_MAX_LENGTH_DEFAULT,
+    VLDataModule,
+)
+from flava.native.model import FLAVAPreTrainModule, get_optimizer
+from flava.native.utils import (
+    build_config,
+    enable_tf32,
+    get_model_parameters,
+    get_model_size_gb,
+    move_to_device,
+    print0,
+    run_imagenet_zero_shot,
+    set_seed,
+    setup_distributed_device,
+)
+from flava.utils import build_datamodule_kwargs
+
+from omegaconf import DictConfig, OmegaConf
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    apply_activation_checkpointing,
+    checkpoint_wrapper,
+    CheckpointImpl,
+)
+from torch.distributed.elastic.multiprocessing.errors import record
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision
+from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler
+from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from torchmultimodal.models.flava.image_encoder import ImageTransformer
+from torchmultimodal.models.flava.text_encoder import BERTTextEncoder
+from torchmultimodal.models.flava.transformer import (
+    FLAVATransformerWithoutEmbeddings,
+    TransformerEncoderLayer,
+)
+from torchmultimodal.modules.losses.flava import FLAVAPretrainingLossOutput
+
+
+def get_datamodules(config: FLAVAArguments) -> Tuple[MultiDataModule, ImageDataModule]:
+    datamodules = []
+
+    # also needed for the imagenet eval callback
+    imagenet_datamodule = ImageDataModule(
+        **build_datamodule_kwargs(config.datasets.image, config.training)
+    )
+    for dataset in config.datasets.selected:
+        if dataset == "image":
+            datamodules.append(imagenet_datamodule)
+        elif dataset == "text":
+            datamodules.append(
+                MLMDataModule(
+                    **build_datamodule_kwargs(config.datasets.text, config.training)
+                )
+            )
+        elif dataset == "vl":
+            datamodules.append(
+                VLDataModule(
+                    **build_datamodule_kwargs(config.datasets.vl, config.training)
+                )
+            )
+        else:
+            raise ValueError(f"unknown dataset: {dataset}")
+
+    return MultiDataModule(datamodules), imagenet_datamodule
+
+
+@record
+class Trainer:
+    def __init__(self, config: DictConfig):
+        if config.training.seed != -1:
+            set_seed(config.training.seed)
+
+        self.device: torch.device = setup_distributed_device()
+        self.config: DictConfig = config
+        self.rank: int = dist.get_rank()
+        self._logger: SummaryWriter = SummaryWriter(
+            f"logs/{config.training.strategy}/{int(time.time())}"
+        )
+        self.steps: int = -1
+        self.epochs: int = -1
+
+        multi_module, image_module = get_datamodules(config)
+
+        self.datamodule: MultiDataModule = multi_module
+        self.datamodule.setup("fit")
+
+        self.imagenet_val_dataloader = image_module.val_dataloader()
+        self.imagenet_val_text_transform = default_text_transform(
+            max_text_length=VL_MAX_LENGTH_DEFAULT
+        )
+
+        self.half_dtype = (
+            torch.bfloat16
+            if config.training.half_precision_format == "bfloat16"
+            else torch.float16
+        )
+
+        self.scaler = ShardedGradScaler() if config.training.enable_amp else None
+
+    def log(
+        self,
+        name: str,
+        value: Union[torch.Tensor, float, int],
+        log_rank_0: bool = True,
+        always_log: bool = False,
+    ):
+        if log_rank_0 and self.rank != 0:
+            return
+
+        if always_log or self.steps % self.config.training.log_interval == 0:
+            self._logger.add_scalar(name, value, self.steps)
+
+    def create_model(self) -> torch.nn.Module:
+        model_config = self.config.get("model", {})
+        print0(f"using model config: {model_config}")
+
+        model = FLAVAPreTrainModule(**model_config)
+        strategy = self.config.training.strategy
+
+        print0(
+            f"before {strategy} model parameters: {get_model_parameters(model):,}, "
+            f"size: {get_model_size_gb(model):.3} GB"
+        )
+
+        if self.config.training.activation_checkpointing:
+            check_fn = lambda submodule: isinstance(submodule, TransformerEncoderLayer)
+            checkpoint_impl = CheckpointImpl.REENTRANT
+
+            # DDP gradient hooks have compatibility issues with REENTRANT autograd
+            if strategy == "ddp":
+                checkpoint_impl = CheckpointImpl.NO_REENTRANT
+
+            checkpoint_wrapper_fn = partial(
+                checkpoint_wrapper,
+                offload_to_cpu=False,
+                checkpoint_impl=checkpoint_impl,
+            )
+            apply_activation_checkpointing(
+                model,
+                checkpoint_wrapper_fn=checkpoint_wrapper_fn,
+                check_fn=check_fn,
+            )
+
+        if strategy == "ddp":
+            # TODO do we have to do this in FSDP too? see https://github.com/pytorch/pytorch/issues/75478
+            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+            model = model.to(self.device)
+
+            print0(
+                f"after moving to cuda: {torch.cuda.memory_allocated()/1024**3:.3} GB"
+            )
+
+            model = DDP(
+                model,
+                device_ids=[self.rank],
+                find_unused_parameters=True,
+                gradient_as_bucket_view=True,
+            )
+            print0(f"after DDP: {torch.cuda.memory_allocated()/1024**3:.3} GB")
+        elif strategy == "fsdp":
+            mp = None
+            if self.config.training.enable_half_reduce_in_fsdp:
+                mp = MixedPrecision(
+                    # param_dtype=self.half_dtype,  not working
+                    reduce_dtype=self.half_dtype,
+                    # buffer_dtype=self.half_dtype,
+                )
+
+            model = FSDP(
+                model,
+                mixed_precision=mp,
+                device_id=self.device,
+                auto_wrap_policy=partial(
+                    transformer_auto_wrap_policy,
+                    transformer_layer_cls={
+                        TransformerEncoderLayer,
+                        ImageTransformer,
+                        BERTTextEncoder,
+                        FLAVATransformerWithoutEmbeddings,
+                    },
+                ),
+                limit_all_gathers=True,
+            )
+
+            print0(f"after FSDP {torch.cuda.memory_allocated()/1024**3:.3} GB")
+
+        else:
+            raise ValueError(f"unknown strategy: {strategy}")
+
+        print0(
+            f"after {strategy} model parameters: {get_model_parameters(model):,}, "
+            f"size: {get_model_size_gb(model):.3} GB"
+        )
+
+        return model
+
+    def calculate_loss(
+        self, output: FLAVAPretrainingLossOutput, validation=False
+    ) -> torch.Tensor:
+        losses = output.losses
+
+        total_loss = 0
+        for key in losses:
+            if losses[key] is not None:
+                total_loss += losses[key]
+                loss_reduce = losses[key].detach()
+                dist.reduce(loss_reduce, dst=0)
+                if validation:
+                    mode = "validation"
+                else:
+                    mode = "train"
+                self.log(
+                    f"{mode}/losses/{key}",
+                    loss_reduce.item() / dist.get_world_size(),
+                )
+
+        return total_loss
+
+    def preprocess_data(self, data: Dict[str, Any]):
+        data = self.datamodule.on_before_batch_transfer(data, None)
+        data = move_to_device(data, self.device)
+        return self.datamodule.on_after_batch_transfer(data, None)
+
+    def _log_iteration_times(self, iteration_times):
+        profile_warmup_steps = config.get("profile_warmup_steps", 100)
+        start_idx = (
+            profile_warmup_steps
+            if profile_warmup_steps < self.config.training.max_steps
+            else 0
+        )
+        iteration_times = iteration_times[start_idx:]
+        avg_it_time = np.mean(iteration_times)
+        avg_throughput = (
+            config.training.batch_size * dist.get_world_size()
+        ) / avg_it_time
+        print0(f"Average over {len(iteration_times)} steps")
+        print0(f"Average iteration time {round(avg_it_time,4)}")
+        print0(f"Average throughput {round(avg_throughput,4)}")
+
+    def train(self) -> None:
+        print0(OmegaConf.to_container(self.config.training))
+        self.model = self.create_model()
+        model = self.model
+
+        optimizer, scheduler = get_optimizer(
+            model,
+            **self.config.training.optimizer,
+        )
+
+        iteration_times = []
+
+        while True:
+            t0 = time.time()
+            self.epochs += 1
+            dataloader = self.datamodule.train_dataloader()
+            dataloader.set_epoch(self.epochs)
+
+            for i, data in enumerate(dataloader):
+                torch.cuda.reset_peak_memory_stats()
+
+                self.steps += 1
+
+                if self.config.training.max_steps < self.steps:
+                    if self.rank == 0:
+                        self._log_iteration_times(iteration_times)
+                    print0("Max steps reached, exiting")
+                    return
+
+                model.train()
+                data = self.preprocess_data(data)
+                optimizer.zero_grad(set_to_none=True)
+
+                with torch.cuda.amp.autocast(
+                    dtype=self.half_dtype, enabled=bool(self.scaler)
+                ):
+                    output = model(data)
+                print0(
+                    f"after forward pass {torch.cuda.memory_allocated()/1024**3:.3} GB"
+                )
+                self.log(
+                    "stats/fwd memory alloc",
+                    torch.cuda.memory_allocated() / 1024**3,
+                )
+                self.log(
+                    "stats/fwd memory reserved",
+                    torch.cuda.memory_reserved() / 1024**3,
+                )
+
+                total_loss = self.calculate_loss(output)
+
+                if self.scaler:
+                    self.scaler.scale(total_loss).backward()
+                    self.scaler.step(optimizer)
+                    self.scaler.update()
+                else:
+                    total_loss.backward()
+                    optimizer.step()
+
+                scheduler.step()
+                torch.cuda.synchronize()
+                t1 = time.time()
+                batch_time = t1 - t0
+                batch_size = config.training.batch_size * dist.get_world_size()
+                items_time = batch_size / (t1 - t0)
+
+                t0 = t1
+                self.log("stats/sec per batch", batch_time)
+                self.log("stats/items per sec", items_time)
+
+                total_loss = total_loss.detach()
+                dist.reduce(total_loss, dst=0)
+
+                if self.rank == 0:
+                    norm_total_loss = total_loss.item() / dist.get_world_size()
+
+                    print(
+                        f"epoch: {self.epochs} step {self.steps} loss: {norm_total_loss:.4}"
+                    )
+                    self.log("train/loss", norm_total_loss)
+                    self.log("stats/batch_size", batch_size)
+
+                    iteration_times.append(batch_time)
+
+                    cuda_info = torch.cuda.memory_stats()
+                    print("cuda alloc retries ", cuda_info.get("num_alloc_retries", 0))
+
+                self.log(
+                    "stats/max_gpu_allocated_gb",
+                    torch.cuda.max_memory_allocated() / 1024**3,
+                )
+                # TODO implement imagenet eval
+                # TODO implement checkpoint saving
+
+                self.validate()
+
+    def validate(self):
+        if self.steps % self.config.training.validation_steps != 0 or self.steps == 0:
+            return
+
+        model = self.model
+        model.eval()
+        print0("evaluating")
+
+        validation_loader = self.datamodule.val_dataloader()
+        validation_loss = torch.Tensor([0]).to(self.device)
+
+        for data in validation_loader:
+            data = self.preprocess_data(data)
+            with torch.no_grad():
+                with torch.cuda.amp.autocast(
+                    dtype=self.half_dtype, enabled=bool(self.scaler)
+                ):
+                    output = model(data)
+                    total_loss = self.calculate_loss(output, validation=True)
+                    validation_loss += total_loss.detach()
+
+        dist.reduce(validation_loss, dst=0)
+        norm_validation_loss = validation_loss.item() / dist.get_world_size()
+
+        print0(f"step {self.steps} EVAL loss: {norm_validation_loss:.4}")
+
+    def imagenet_validate(self):
+        print0("imagenet validation")
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(
+                dtype=self.half_dtype, enabled=bool(self.scaler)
+            ):
+                metrics = run_imagenet_zero_shot(
+                    self.model,
+                    self.imagenet_val_dataloader,
+                    self.device,
+                    self.imagenet_val_text_transform,
+                )
+                if metrics is not None:
+                    for key in metrics:
+                        self.log(
+                            f"val/imagenet/{key}",
+                            metrics[key],
+                            always_log=True,
+                        )
+
+
+if __name__ == "__main__":
+    datasets.logging.set_verbosity_error()  # too spammy
+
+    config: FLAVAArguments = build_config()
+    if config.training.enable_tf32:
+        enable_tf32()
+
+    trainer = Trainer(config)
+    trainer.train()

multimodal/examples/flava/native/utils.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+from typing import Any
+
+import torch
+from flava.data.imagenet_zeroshot_data import (
+    imagenet_classnames,
+    openai_imagenet_template,
+)
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+from torch import distributed as dist
+from tqdm import tqdm
+
+# optional syntax-highlighting for console output
+try:
+    from rich.console import Console
+
+    c = Console(force_terminal=True)
+    print = c.log
+except ImportError:
+    pass
+
+
+def build_config() -> DictConfig:
+    cli_conf = OmegaConf.from_cli()
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = instantiate(yaml_conf)
+    conf = OmegaConf.merge(conf, cli_conf)
+    return conf
+
+
+# TODO replace with tlc.copy_data_to_device
+def move_to_device(obj: Any, device: torch.device) -> Any:
+    if isinstance(obj, dict):
+        d = {}
+        for k, v in obj.items():
+            d[k] = move_to_device(v, device)
+        return d
+    if isinstance(obj, list):
+        l = []
+        for v in obj:
+            l.append(move_to_device(v, device))
+        return l
+
+    return obj.to(device)
+
+
+def get_model_size_gb(model: torch.nn.Module) -> int:
+    return sum(p.numel() * p.element_size() for p in model.parameters()) / (1024**3)
+
+
+def get_model_parameters(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+def set_seed(seed: int) -> None:
+    torch.manual_seed(seed)
+    random.seed(seed)
+
+
+def setup_distributed_device() -> torch.device:
+    if not torch.cuda.is_available() or not dist.is_available():
+        return torch.device("cpu")
+
+    dist.init_process_group("nccl")
+    local_rank = int(os.environ["LOCAL_RANK"])
+    print("local rank", local_rank)
+    torch.cuda.set_device(local_rank)
+    return torch.device(f"cuda:{local_rank}")
+
+
+def print0(*args, **kwargs) -> None:
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        print(*args, **kwargs)
+
+
+def enable_tf32() -> None:
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+
+def rank0_only(func):
+    def wrapper(*args, **kwargs):
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+# zero shot classifier functions
+
+
+def _zero_shot_classifier(model, device, text_transform, *args, **kwargs):
+    zeroshot_weights = []
+    for classname in tqdm(imagenet_classnames):
+        texts = text_transform(
+            [template(classname) for template in openai_imagenet_template]
+        )["input_ids"]
+        texts = texts.to(device)
+        class_embeddings = model(texts, action="encode_text")
+        class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
+        class_embedding = class_embeddings.mean(dim=0)
+        class_embedding /= class_embedding.norm()
+        zeroshot_weights.append(class_embedding)
+
+    zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
+    return zeroshot_weights
+
+
+def _accuracy(output, target, topk=(1,)):
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    return [
+        float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy())
+        for k in topk
+    ]
+
+
+def run_imagenet_zero_shot(model, dataloader, device, text_transform, *args, **kwargs):
+    print0("Starting ImageNet Zero-Shot Eval")
+    print0("Building classifier")
+    classifier = _zero_shot_classifier(model, device, text_transform)
+    print0("Classifier built")
+    top1, top5, n = 0.0, 0.0, 0.0
+    for i, sample in tqdm(enumerate(dataloader)):
+        images = sample["image"]
+        target = sample["label"]
+        images = images.to(device)
+        target = target.to(device)
+
+        # predict
+        # if hasattr(model, "module"):
+        #     image_features = model.module.encode_image({"image": images})
+        # else:
+        image_features = model(images, action="encode_image")
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        logits = 100.0 * image_features @ classifier
+
+        # measure accuracy
+        acc1, acc5 = _accuracy(logits, target, topk=(1, 5))
+        top1 += acc1
+        top5 += acc5
+        n += images.size(0)
+        if i == 5:
+            break
+
+    top1 = top1 / n
+    top5 = top5 / n
+    results = {}
+    results["imagenet-zeroshot-val-top1"] = top1
+    results["imagenet-zeroshot-val-top5"] = top5
+    print0("results: ", results)
+    return results

multimodal/examples/flava/notebooks/RemapFLAVACheckpoint.ipynb

@@ -0,0 +1,172 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "7cc982d1",
+   "metadata": {},
+   "source": [
+    "# Re-map FLAVA checkpoint\n",
+    "\n",
+    "Modifying FLAVA's components can cause existing model checkpoints to go out of sync with the updated architecture. This notebook shows how to load the existing checkpoint, re-map the old layers to the new layers, and save the new checkpoint.\n",
+    "\n",
+    "To upload a new checkpoint, you must have access to the PyTorch AWS S3 account, and manually upload it from a local copy."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "411e4191",
+   "metadata": {},
+   "source": [
+    "### Load original model\n",
+    "\n",
+    "Load the existing checkpoint into the FLAVA class to see what the architecture currently is."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "88ee917b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from torchmultimodal.models.flava.model import flava_model_for_classification, flava_model_for_pretraining\n",
+    "\n",
+    "# flava_classification = flava_model_for_classification(num_classes=3)\n",
+    "flava_pretraining = flava_model_for_pretraining(pretrained_model_key='flava_full')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5f00b369",
+   "metadata": {},
+   "source": [
+    "### Print summary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cc286394",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flava_pretraining"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d774455",
+   "metadata": {},
+   "source": [
+    "### Mapping function\n",
+    "\n",
+    "Replace this function with the code needed to map the old layer weights to the new layer weights."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "cc9e4537",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "\n",
+    "def map_state_dict(state_dict):\n",
+    "    mapped_state_dict = {}\n",
+    "    for param, val in state_dict.items():\n",
+    "        res = re.search('attention.attention', param)\n",
+    "        if res:\n",
+    "            idx = res.start()\n",
+    "            new_param = param[:idx] + param[idx+10:]\n",
+    "        else:\n",
+    "            new_param = param\n",
+    "        mapped_state_dict[new_param] = val\n",
+    "    return mapped_state_dict"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "29870590",
+   "metadata": {},
+   "source": [
+    "### Load old state dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "41f64d26",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load from url, replace this path if it changes\n",
+    "# old_model_url = 'https://download.pytorch.org/models/multimodal/flava/flava_model.pt'\n",
+    "# old_state_dict = torch.hub.load_state_dict_from_url(old_model_url)\n",
+    "\n",
+    "# Or get from loaded model\n",
+    "old_state_dict = flava_pretraining.model.state_dict()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "75322113",
+   "metadata": {},
+   "source": [
+    "### Perform re-mapping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "17363ae8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#new_state_dict = map_state_dict(old_state_dict)\n",
+    "new_state_dict = old_state_dict"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d94c4133",
+   "metadata": {},
+   "source": [
+    "### Save updated checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "bc6baad9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_path = '/Users/rafiayub/flava_model.pt'\n",
+    "torch.save(new_state_dict, save_path)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

multimodal/examples/flava/tools/convert_weights.py

@@ -0,0 +1,72 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+
+import torch
+from torchmultimodal.models.flava.model import flava_model_for_pretraining
+
+KEY_REPLACEMENTS = {
+    "image_encoder.module": "image_encoder",
+    "text_encoder.module": "text_encoder",
+    "mm_encoder.module": "mm_encoder",
+    "mm_encoder.encoder.cls_token": "mm_encoder.cls_token",
+    "mm_image_projection": "image_to_mm_projection",
+    "mm_text_projection": "text_to_mm_projection",
+    "model.heads.cmd.mim_head": "loss.mmm_loss.mim",
+    "model.heads.cmd.mlm_head": "loss.mmm_loss.mlm",
+    "model.heads.fairseq_mlm": "loss.mlm_loss",
+    "model.heads.imagenet.mim_head": "loss.mim_loss",
+    "cls.predictions.transform": "cls",
+    "cls.predictions": "cls",
+    "cls.LayerNorm": "cls.layer_norm",
+    "model.text_projection": "loss.contrastive_loss.text_projection",
+    "model.image_projection": "loss.contrastive_loss.image_projection",
+    "model.heads.cmd.clip_head.logit_scale": "loss.contrastive_loss.logit_scale",
+    "model.heads.cmd.itm_head": "loss.itm_loss",
+    "intermediate.dense": "intermediate",
+    "output.dense": "output",
+}
+
+
+def convert_weights(args):
+    ckpt = torch.load(args.ckpt_file, map_location="cpu")
+    flava = flava_model_for_pretraining()
+    model = ckpt["model"]
+    import pdb
+
+    pdb.set_trace()
+    for key in list(model.keys()):
+        original = key
+        for option, replacement in KEY_REPLACEMENTS.items():
+            key = key.replace(option, replacement)
+        model[key] = model.pop(original)
+
+    if args.add_codebook:
+        # Since codebook is anyways not trained in FLAVA pretraining
+        # we can use the pretrained one that we get from FLAVA initialized
+        # model
+        model.update(
+            {
+                f"image_codebook.{key}": value
+                for key, value in flava.image_codebook.state_dict().items()
+            }
+        )
+    flava.load_state_dict(model)
+
+    # Let's save the model now.
+    torch.save(flava.state_dict(), args.save_file)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert weights")
+    parser.add_argument("ckpt_file", type=str)
+    parser.add_argument("save_file", type=str)
+    parser.add_argument("--add_codebook", action="store_true")
+
+    args = parser.parse_args()
+
+    convert_weights(args)

multimodal/examples/mugen/data/README.md

@@ -0,0 +1,10 @@
+This folder contains code for interfacing the [MUGEN dataset](https://mugen-org.github.io). The MUGEN dataset contains over 300k videos, each with corresponding audio and text, from the game CoinRun.
+
+Before using this code,
+
+1. Download the 3.2s-video dataset [here](https://mugen-org.github.io/download) and save as `datasets/coinrun` in your working directory.
+    * In each of `datasets/coinrun/coinrun_dataset_jsons/release/{train/val/test}.json`, change the value of `json_object["metadata"]["data_folder"]` to the absolute path of `datasets/coinrun`, e.g. `"/path/to/datasets/coinrun/"`.
+2. Download the MUGEN dataset assets [here](https://github.com/mugen-org/MUGEN_baseline/tree/main/lib/data/coinrun/assets) and save under `datasets/coinrun` as `datasets/coinrun/assets` in your pwd.
+    * Downloading the assets from GitHub requires `git clone`-ing the original MUGEN repo and copying the assets directory located at `MUGEN_baseline/lib/data/coinrun/assets`.
+
+Note: saving the dataset and assets to locations other than those listed above requires passing custom arguments to `MUGENDataModuleBase` or `MUGENDataset` through `MUGENDatasetArgs.data_path` and `MUGENDatasetArgs.asset_path`, respectively.

multimodal/examples/mugen/data/coinrun/construct_from_json.py

@@ -0,0 +1,756 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import os
+
+import numpy as np
+from PIL import Image
+
+DEATH_ANIM_LENGTH = 30
+FINISHED_LEVEL_ANIM_LENGTH = 20
+MONSTER_DEATH_ANIM_LENGTH = 3
+SPACE = "."
+LADDER = "="
+LAVA_SURFACE = "^"
+LAVA_MIDDLE = "|"
+WALL_SURFACE = "S"
+WALL_MIDDLE = "A"
+WALL_CLIFF_LEFT = "a"
+WALL_CLIFF_RIGHT = "b"
+COIN_OBJ1 = "1"
+COIN_OBJ2 = "2"
+CRATE_NORMAL = "#"
+CRATE_DOUBLE = "$"
+CRATE_SINGLE = "&"
+CRATE_WARNING = "%"
+
+
+def define_semantic_color_map(max_label=18):
+    assert max_label in [18, 21, 22], f"max_label {max_label} is not supported!"
+
+    semantic_color_map = {}
+
+    semantic_color_map["background"] = 0
+
+    # alien is always set to max_label (assumes it always appear in a video)
+    semantic_color_map["alien"] = max_label
+
+    if max_label == 18:
+        semantic_color_map["world"] = {
+            WALL_MIDDLE: 3,
+            WALL_SURFACE: 4,
+            WALL_CLIFF_LEFT: 5,
+            WALL_CLIFF_RIGHT: 6,
+            COIN_OBJ1: 17,
+            COIN_OBJ2: 0,
+            CRATE_NORMAL: 8,
+            CRATE_DOUBLE: 8,
+            CRATE_SINGLE: 8,
+            CRATE_WARNING: 8,
+            LAVA_MIDDLE: 1,
+            LAVA_SURFACE: 2,
+            LADDER: 7,
+        }
+
+        semantic_color_map["shield"] = 0
+
+        semantic_color_map["monster"] = {
+            "sawHalf": 16,
+            "bee": 15,
+            "slimeBlock": 14,
+            "slimeBlue": 13,
+            "mouse": 12,
+            "snail": 11,
+            "ladybug": 10,
+            "wormPink": 9,
+            "barnacle": 0,
+            "frog": 0,
+        }
+    else:
+        semantic_color_map["world"] = {
+            WALL_MIDDLE: 3,
+            WALL_SURFACE: 4,
+            WALL_CLIFF_LEFT: 5,
+            WALL_CLIFF_RIGHT: 6,
+            COIN_OBJ1: 19,
+            COIN_OBJ2: 20,
+            CRATE_NORMAL: 8,
+            CRATE_DOUBLE: 8,
+            CRATE_SINGLE: 8,
+            CRATE_WARNING: 8,
+            LAVA_MIDDLE: 1,
+            LAVA_SURFACE: 2,
+            LADDER: 7,
+        }
+
+        semantic_color_map["shield"] = 21
+
+        semantic_color_map["monster"] = {
+            "sawHalf": 16,
+            "bee": 15,
+            "slimeBlock": 14,
+            "slimeBlue": 13,
+            "mouse": 12,
+            "snail": 11,
+            "ladybug": 10,
+            "wormPink": 9,
+            "barnacle": 17,
+            "frog": 18,
+        }
+
+    return semantic_color_map
+
+
+def generate_asset_paths(game):
+    # use background corresponding with ground theme
+    bgtheme = game.background_themes[game.world_theme_n]
+
+    gtheme = game.ground_themes[game.world_theme_n]
+    walls = "kenney/Ground/" + gtheme + "/" + gtheme.lower()
+
+    # default option with fixed agent look
+    atheme = game.agent_themes[game.agent_theme_n]
+    alien = "kenneyLarge/Players/128x256_no_helmet/" + atheme + "/alien" + atheme
+    alien_paths = {"Mugen": alien}
+
+    tiles = "kenney/Tiles/"
+    items = "kenneyLarge/Items/"
+    enemy = "kenneyLarge/Enemies/"
+
+    asset_files = {}
+
+    asset_files["background"] = bgtheme
+
+    asset_files["world"] = {
+        WALL_MIDDLE: walls + "Center.png",
+        WALL_SURFACE: walls + "Mid.png",
+        WALL_CLIFF_LEFT: walls + "Cliff_left.png",
+        WALL_CLIFF_RIGHT: walls + "Cliff_right.png",
+        COIN_OBJ1: items + "coinGold.png",
+        COIN_OBJ2: items + "gemRed.png",
+        CRATE_NORMAL: tiles + "boxCrate.png",
+        CRATE_DOUBLE: tiles + "boxCrate_double.png",
+        CRATE_SINGLE: tiles + "boxCrate_single.png",
+        CRATE_WARNING: tiles + "boxCrate_warning.png",
+        LAVA_MIDDLE: tiles + "lava.png",
+        LAVA_SURFACE: tiles + "lavaTop_low.png",
+        LADDER: tiles + "ladderMid.png",
+    }
+
+    asset_files["alien"] = {}
+    for alien_name in alien_paths.keys():
+        asset_files["alien"][alien_name] = {
+            "walk1": alien_paths[alien_name] + "_walk1.png",
+            "walk2": alien_paths[alien_name] + "_walk2.png",
+            "climb1": alien_paths[alien_name] + "_climb1.png",
+            "climb2": alien_paths[alien_name] + "_climb2.png",
+            "stand": alien_paths[alien_name] + "_stand.png",
+            "jump": alien_paths[alien_name] + "_jump.png",
+            "duck": alien_paths[alien_name] + "_duck.png",
+            "hit": alien_paths[alien_name] + "_hit.png",
+        }
+    asset_files["shield"] = "bubble_shield.png"
+
+    game.flatten_monster_names()
+    # monster assets are generated based on list of names used at rendering
+    asset_files["monster"] = {
+        name: enemy + name + ".png" for name in game.flattened_monster_names
+    }
+
+    return asset_files
+
+
+# binarize alpha channel if input img is in RGBA mode, set anything above 0 to 255
+def binarize_alpha_channel(img):
+    if img.mode != "RGBA":
+        return img
+
+    w, h = img.size
+    for i in range(w):
+        for j in range(h):
+            pixel = img.getpixel((i, j))
+
+            # set alpha to 255 if alpha > 0
+            if pixel[3] > 0:
+                img.putpixel((i, j), (pixel[0], pixel[1], pixel[2], 255))
+
+    return img
+
+
+class Asset:
+    def __init__(
+        self,
+        name,
+        file,
+        asset_root,
+        kind="world",
+        kx=80,
+        ky=80,
+        semantic_color=(0, 0, 0),
+        flip=False,
+        binarize_alpha=False,
+    ):
+        self.name = name
+        self.file = file
+        self.asset_root = asset_root
+        self.kind = kind
+        self.kx = kx
+        self.ky = ky
+        self.semantic_color = semantic_color
+        self.flip = flip
+        self.binarize_alpha = binarize_alpha
+
+        self.load_asset()
+
+    def load_asset(self):
+        asset_path = os.path.join(self.asset_root, self.file)
+        if not os.path.isfile(asset_path):
+            # basically remove the '_walk1' postfix
+            fallback_path = (
+                "_".join(asset_path.split("_")[:-1]) + "." + asset_path.split(".")[-1]
+            )
+            assert os.path.isfile(fallback_path), asset_path
+            asset_path = fallback_path
+        self.asset = Image.open(asset_path)
+
+        # used for (user control) asset swap, because alien h:w == 2:1 while others is 1:1
+        # the asset resize at loading and render grid size all need to change respectively
+        self.aspect_ratio = self.asset.size[1] / self.asset.size[0]
+
+        if self.kind == "world":
+            if self.name != LAVA_MIDDLE and self.name != LAVA_SURFACE:
+                # LAVA has a special way of rendering animation so don't resize now
+                self.asset = self.asset.resize(
+                    (math.ceil(self.kx + 0.5), math.ceil(self.ky + 0.5))
+                )
+        elif self.kind == "alien":
+            self.asset = self.asset.resize(
+                (math.ceil(self.kx), math.ceil(self.aspect_ratio * self.ky))
+            )
+        elif self.kind == "shield":
+            self.asset = self.asset.resize(
+                (math.ceil(self.kx * 1.15), math.ceil(self.ky * 2.1))
+            )
+        elif self.kind == "monster" or self.kind == "background":
+            self.asset = self.asset.resize((math.ceil(self.kx), math.ceil(self.ky)))
+        else:
+            raise NotImplementedError(f"Unknown asset kind {self.kind}")
+
+        # flip if needed (for facing left/right)
+        if self.flip:
+            self.asset = self.asset.transpose(Image.FLIP_LEFT_RIGHT)
+
+        if self.binarize_alpha:
+            self.asset = binarize_alpha_channel(self.asset)
+
+
+def load_assets(
+    asset_files, asset_root, semantic_color_map, kx=80, ky=80, gen_original=False
+):
+    asset_map = {}
+
+    for kind in asset_files.keys():
+        assert kind in semantic_color_map
+
+        if kind == "background":
+            # background will be loaded separately
+            continue
+
+        if kind == "shield":
+            # asset file for the bubble shield in agent power-up mode
+            asset_map[kind] = Asset(
+                name=kind,
+                file=asset_files[kind],
+                asset_root=asset_root,
+                kind=kind,
+                kx=kx,
+                ky=ky,
+                semantic_color=semantic_color_map[kind],
+                binarize_alpha=not gen_original,
+            )
+            continue
+
+        for key in asset_files[kind].keys():
+            if kind == "world":
+                # ground asset, no need to worry about pose or facing
+                asset_map[key] = Asset(
+                    name=key,
+                    file=asset_files[kind][key],
+                    asset_root=asset_root,
+                    kind=kind,
+                    kx=kx,
+                    ky=ky,
+                    semantic_color=semantic_color_map[kind][key],
+                    binarize_alpha=not gen_original,
+                )
+            elif kind == "alien":
+                for pose in asset_files[kind][key].keys():
+                    # facing right is default to empty
+                    all_facings = ["", "_left"]
+                    for facing in all_facings:
+                        a_key = key + "_" + pose + facing
+
+                        asset_map[a_key] = Asset(
+                            name=a_key,
+                            file=asset_files[kind][key][pose],
+                            asset_root=asset_root,
+                            kind=kind,
+                            kx=kx,
+                            ky=ky,
+                            semantic_color=semantic_color_map[kind],
+                            flip=(facing != ""),  # flip the asset if facing is not ''
+                            binarize_alpha=not gen_original,
+                        )
+            elif kind == "monster":
+                # for monsters, 3 types of assets will be loaded
+                # for each of them, facing can be left or right
+                all_poses = ["", "_move", "_dead"]  # walk1 is default to empty
+                all_facings = ["", "_right"]  # facing left is default to empty
+                base_fn = os.path.splitext(asset_files[kind][key])[
+                    0
+                ]  # e.g. Enemies/bee
+                for pose in all_poses:
+                    for facing in all_facings:
+                        m_key = key + pose + facing
+                        file_name = base_fn + pose + ".png"
+
+                        asset_map[m_key] = Asset(
+                            name=m_key,
+                            file=file_name,
+                            asset_root=asset_root,
+                            kind="monster",
+                            kx=kx,
+                            ky=ky,
+                            semantic_color=semantic_color_map[kind][key],
+                            flip=(facing != ""),  # flip the asset if facing is not ''
+                            binarize_alpha=not gen_original,
+                        )
+            else:
+                raise NotImplementedError(f"Unknown asset kind {kind}")
+
+    return asset_map
+
+
+# load background asset, zoom is different so need a separate function
+def load_bg_asset(asset_files, asset_root, semantic_color_map, zx, zy):
+    kind = "background"
+    bg_asset = Asset(
+        name=kind,
+        file=asset_files[kind],
+        asset_root=asset_root,
+        kind=kind,
+        kx=zx,
+        ky=zy,
+        semantic_color=semantic_color_map[kind],
+    )
+    return bg_asset
+
+
+# used for alien dying animation in gen_original mode
+def get_transparent_asset(input_asset, transparency):
+    assert input_asset.mode == "RGBA"
+    np_asset = np.array(input_asset, dtype=np.int16)
+    np_asset[:, :, 3] -= transparency
+    np_asset[:, :, 3] = np.clip(np_asset[:, :, 3], 0, None)
+    return Image.fromarray(np_asset.astype(np.uint8))
+
+
+# return rect in integer values, floor for x1,y1, ceil for x2,y2 or w,h
+def integer_rect(rect):
+    return [
+        math.floor(rect[0]),
+        math.floor(rect[1]),
+        math.ceil(rect[2]),
+        math.ceil(rect[3]),
+    ]
+
+
+def convert_xywh_to_xyxy(rect):
+    return [rect[0], rect[1], rect[0] + rect[2], rect[1] + rect[3]]
+
+
+def convert_xyxy_to_xywh(rect):
+    return [rect[0], rect[1], rect[2] - rect[0], rect[3] - rect[1]]
+
+
+# rect format is xywh, img_size is (w,h)
+def check_out_of_bounds(rect, img_size):
+    if rect[0] + rect[2] < 0:
+        return True
+    if rect[0] > img_size[0]:
+        return True
+    if rect[1] + rect[3] < 0:
+        return True
+    if rect[1] > img_size[1]:
+        return True
+    return False
+
+
+# return intersect of two rects, input and output are both in xywh format
+def intersect_rects(rect1, rect2):
+    xyxy_rect1 = convert_xywh_to_xyxy(rect1)
+    xyxy_rect2 = convert_xywh_to_xyxy(rect2)
+    xyxy_res_rect = [
+        max(xyxy_rect1[0], xyxy_rect2[0]),
+        max(xyxy_rect1[1], xyxy_rect2[1]),
+        min(xyxy_rect1[2], xyxy_rect2[2]),
+        min(xyxy_rect1[3], xyxy_rect2[3]),
+    ]
+
+    xywh_res_rect = convert_xyxy_to_xywh(xyxy_res_rect)
+
+    # check if the intersection is empty
+    if xywh_res_rect[2] > 0 and xywh_res_rect[3] > 0:
+        return xywh_res_rect
+    else:
+        return None
+
+
+# rect is in the format of xywh
+def paint_color_in_rect_with_mask(
+    img, rect, color, mask, gen_original=False, ignore_mask=False, cut_mask_top_ratio=0
+):
+    w, h = mask.size
+    img_w, img_h = img.size
+    # in some cases, mask size doesn't match the rect (e.g. monster dying)
+    if rect[2] != w or rect[3] != h:
+        if not gen_original:
+            mask = mask.resize((rect[2], rect[3]), resample=Image.NEAREST)
+        else:
+            mask = mask.resize((rect[2], rect[3]))
+        w, h = mask.size
+
+    if not gen_original:
+        # generate semantic map
+        if ignore_mask and cut_mask_top_ratio != 0:
+            # specifically for agent because its asset has a large empty area in the top,
+            # we don't want it to be fully masked
+            if cut_mask_top_ratio < 0:
+                # automatic calculate the first non-empty row from top
+                np_mask = np.array(mask)
+                cut_mask_top_rows = (np_mask.T[0].sum(axis=0) != 0).argmax(axis=0)
+            else:
+                cut_mask_top_rows = int(cut_mask_top_ratio * rect[2])
+            rect[1] += cut_mask_top_rows
+            rect[3] = mask.size[1] - cut_mask_top_rows
+
+            img = img.paste(color, convert_xywh_to_xyxy(rect))
+        else:
+            # paste in single color if generating semantic maps (so not original)
+            # if ignore_mask, this will generate a complete block mask same as rect
+            img = img.paste(
+                color,
+                convert_xywh_to_xyxy(rect),
+                mask if (mask.mode == "RGBA" and not ignore_mask) else None,
+            )
+    else:
+        # generate rgb data
+        img = img.paste(
+            mask, convert_xywh_to_xyxy(rect), mask if mask.mode == "RGBA" else None
+        )
+
+    return
+
+
+def draw_game_frame(
+    game,
+    frame_id,
+    asset_map,
+    kx,
+    ky,
+    gen_original=False,
+    bbox_smap_for_agent=False,
+    bbox_smap_for_monsters=False,
+    alien_name=None,
+    skip_foreground=False,
+    skip_background=False,
+    skip_mugen=False,
+    only_mugen=False,
+):
+    # set default alien name/key
+    if alien_name is None:
+        alien_name = "Mugen"
+
+    # initialize an empty image (all zero, for background)
+    if not gen_original:
+        img = Image.new("L", (game.video_res, game.video_res))
+    else:
+        img = Image.new("RGB", (game.video_res, game.video_res))
+
+    video_center = (game.video_res - 1) // 2
+
+    frame = game.frames[frame_id]
+
+    # for agent-centric
+    # dx = -frame.agent.x * kx + video_center - 0.5 * kx
+    # dy = frame.agent.y * ky - video_center - 0.5 * ky
+    # for video data (no vertical camera move)
+    dx = -frame.agent.x * kx + video_center - 0.5 * kx
+
+    # different dy/ky ratio based on zoom level, to adjust camera view
+    if game.zoom == 5.5:
+        dy_ratio = 5.0
+    elif game.zoom == 4.3:
+        dy_ratio = 6.5
+    elif game.zoom == 5.0:
+        dy_ratio = 5.5
+    elif game.zoom == 6.0:
+        dy_ratio = 4.5
+    else:
+        raise NotImplementedError(f"zoom level {game.zoom} is not supported!")
+    dy = -video_center + dy_ratio * ky
+
+    # update background image with proper zoom for gen_original mode
+    # NOTE: if desired background label is not zero, set it here to asset_map['background'].semantic_color
+    if gen_original and not skip_background and not only_mugen:
+        zx = game.video_res * game.zoom
+        zy = zx
+        for tile_x in range(-1, 3):
+            for tile_y in range(-1, 2):
+                bg_rect = [0, 0, zx, zy]
+                bg_rect[0] = (
+                    zx * tile_x
+                    + video_center
+                    + game.bgzoom * (dx + kx * game.maze_h / 2)
+                    - zx * 0.5
+                )
+                bg_rect[1] = (
+                    zy * tile_y
+                    + video_center
+                    + game.bgzoom * (dy - ky * game.maze_h / 2)
+                    - zy * 0.5
+                )
+                if check_out_of_bounds(bg_rect, img.size):
+                    continue
+                img.paste(
+                    asset_map["background"].asset,
+                    convert_xywh_to_xyxy(integer_rect(bg_rect)),
+                )
+
+    # NOTE: game engine now hard-code 64 for maze_size
+    radius = int(1 + game.maze_w / game.zoom)
+    ix = int(frame.agent.x + 0.5)
+    iy = int(frame.agent.y + 0.5)
+    x_start = max(ix - radius, 0)
+    x_end = min(ix + radius + 1, game.maze_w)
+    y_start = max(iy - radius, 0)
+    y_end = min(iy + radius + 1, game.maze_h)
+    win_h = game.video_res
+
+    # convert eaten coins to a set for faster checking coordinates
+    coins_eaten_set = {tuple(coin_coord) for coin_coord in frame.coins_eaten}
+
+    if not skip_background and not only_mugen:
+        for y in range(y_start, y_end):
+            for x in range(x_start, x_end):
+                wkey = game.maze[y][x]
+                if wkey == SPACE:
+                    continue
+
+                # eaten coins is treated the same as SPACE, just continue
+                # but we should not modify the coins in maze to SPACE, or it may cause inconsistency
+                # if we ever need to render backwards or save json after drawing
+                if (x, y) in coins_eaten_set:
+                    continue
+
+                assert wkey in asset_map, f"{wkey} not in assets!"
+
+                tile_rect = [
+                    kx * x + dx - 0.1,
+                    win_h - ky * y + dy - 0.1,
+                    kx + 0.5 + 0.2,
+                    ky + 0.5 + 0.2,
+                ]
+
+                # skip tile if the rect is completely out-of-bounds
+                if check_out_of_bounds(tile_rect, img.size):
+                    continue
+
+                if wkey == LAVA_MIDDLE or wkey == LAVA_SURFACE:
+                    d1 = tile_rect[:]
+                    d2 = tile_rect[:]
+                    asset_size = asset_map[wkey].asset.size
+                    sr = [0, 0, asset_size[0], asset_size[1]]
+                    sr1 = sr[:]
+                    sr2 = sr[:]
+                    tr = frame.state_time * 0.1
+                    tr -= int(tr)
+                    tr *= -1
+                    d1[0] += tr * tile_rect[2]
+                    d2[0] += tile_rect[2] + tr * tile_rect[2]
+                    sr1[0] += -tr * asset_size[0]
+                    sr2[0] += -asset_size[0] - tr * asset_size[0]
+                    d1 = intersect_rects(d1, tile_rect)
+                    d2 = intersect_rects(d2, tile_rect)
+                    if d1 is not None:
+                        d1[2] += 0.5
+                    if d2 is not None:
+                        d2[0] -= 0.5
+                        d2[2] += 0.5
+                    sr1 = intersect_rects(sr1, sr)
+                    sr2 = intersect_rects(sr2, sr)
+                    if sr1 is not None and d1 is not None:
+                        # crop and render one half of the asset
+                        crop_mask = asset_map[wkey].asset.crop(
+                            integer_rect(convert_xywh_to_xyxy(sr1))
+                        )
+                        paint_color_in_rect_with_mask(
+                            img,
+                            integer_rect(d1),
+                            asset_map[wkey].semantic_color,
+                            crop_mask,
+                            gen_original=gen_original,
+                        )
+                    if sr2 is not None and d2 is not None:
+                        # crop and render the other half of the asset (swapped places horizontally)
+                        crop_mask = asset_map[wkey].asset.crop(
+                            integer_rect(convert_xywh_to_xyxy(sr2))
+                        )
+                        paint_color_in_rect_with_mask(
+                            img,
+                            integer_rect(d2),
+                            asset_map[wkey].semantic_color,
+                            crop_mask,
+                            gen_original=gen_original,
+                        )
+                else:
+                    paint_color_in_rect_with_mask(
+                        img,
+                        integer_rect(tile_rect),
+                        asset_map[wkey].semantic_color,
+                        asset_map[wkey].asset,
+                        gen_original=gen_original,
+                    )
+
+    if not skip_foreground:
+        if not only_mugen:
+            # paint monsters
+            for mi in range(len(frame.monsters)):
+                if frame.monsters[mi].is_dead:
+                    dying_frame_cnt = max(0, frame.monsters[mi].monster_dying_frame_cnt)
+                    monster_shrinkage = (
+                        (MONSTER_DEATH_ANIM_LENGTH - dying_frame_cnt)
+                        * 0.8
+                        / MONSTER_DEATH_ANIM_LENGTH
+                    )
+                    monster_rect = [
+                        math.floor(kx * frame.monsters[mi].x + dx),
+                        math.floor(
+                            win_h
+                            - ky * frame.monsters[mi].y
+                            + dy
+                            + ky * monster_shrinkage
+                        ),
+                        math.ceil(kx),
+                        math.ceil(ky * (1 - monster_shrinkage)),
+                    ]
+                else:
+                    monster_rect = [
+                        math.floor(kx * frame.monsters[mi].x + dx),
+                        math.floor(win_h - ky * frame.monsters[mi].y + dy),
+                        math.ceil(kx),
+                        math.ceil(ky),
+                    ]
+
+                m_name = game.flattened_monster_names[frame.monsters[mi].theme]
+                # add pose and facing to the key to find correct asset
+                m_pose = "" if frame.monsters[mi].walk1_mode else "_move"
+                if frame.monsters[mi].is_dead:
+                    m_pose = "_dead"
+                m_key = (
+                    m_name + m_pose + ("_right" if frame.monsters[mi].vx > 0 else "")
+                )
+
+                paint_color_in_rect_with_mask(
+                    img,
+                    monster_rect,
+                    asset_map[m_key].semantic_color,
+                    asset_map[m_key].asset,
+                    gen_original=gen_original,
+                    ignore_mask=bbox_smap_for_monsters,
+                )
+
+        if not skip_mugen:
+            # paint agent - do it after monsters so agent is always in front
+            a_key = (
+                alien_name
+                + "_"
+                + frame.agent.pose
+                + ("" if frame.agent.is_facing_right else "_left")
+            )
+            # note how aspect_ratio is used for alien rect, this can be applied to
+            # monster rect to support asset that's not 1:1 (e.g. use alien as monster)
+            alien_rect = [
+                math.floor(kx * frame.agent.x + dx),
+                # math.floor(win_h - ky * (frame.agent.y + 1) + dy),    # default for 2:1 alien, no asset swap
+                math.floor(
+                    win_h
+                    - ky * (frame.agent.y + asset_map[a_key].aspect_ratio - 1)
+                    + dy
+                ),
+                math.ceil(kx),
+                # math.ceil(2 * ky),    # default for 2:1 alien, no asset swap
+                math.ceil(asset_map[a_key].aspect_ratio * ky),
+            ]
+            if frame.agent.is_killed:
+                transparency = (
+                    DEATH_ANIM_LENGTH + 1 - frame.agent.killed_animation_frame_cnt
+                ) * 12
+                # only render if not fully transparent
+                if transparency > 255:
+                    agent_asset = None
+                else:
+                    if gen_original:
+                        agent_asset = get_transparent_asset(
+                            asset_map[a_key].asset, transparency
+                        )
+                    else:
+                        # when generating semantic map, alien mask won't change unless fully transparent
+                        agent_asset = asset_map[a_key].asset
+            else:
+                agent_asset = asset_map[a_key].asset
+            if agent_asset is not None:
+                paint_color_in_rect_with_mask(
+                    img,
+                    alien_rect,
+                    asset_map[a_key].semantic_color,
+                    agent_asset,
+                    gen_original=gen_original,
+                    ignore_mask=bbox_smap_for_agent,
+                    cut_mask_top_ratio=0.8,
+                )
+
+            # paint the bubble shield if agent is in power-up mode
+            if frame.agent.power_up_mode:
+                shield_rect = [
+                    # NOTE: game engine hard-codes 7 and 8 for co-ordinates which won't work with video-res that's not 1024
+                    # (for training we usually generate with 256 or 128 video_res), so need to convert them
+                    math.floor(kx * frame.agent.x + dx - 7 * game.video_res / 1024),
+                    math.floor(
+                        win_h
+                        - ky * (frame.agent.y + 1)
+                        + dy
+                        + 8 * game.video_res / 1024
+                    ),
+                    math.ceil(kx * 1.15),
+                    math.ceil(ky * 2.1),
+                ]
+                # pull bubble down when Mugen crouches
+                if frame.agent.pose == "duck":
+                    shield_rect[1] += math.floor(8 * game.video_res / 1024)
+
+                paint_color_in_rect_with_mask(
+                    img,
+                    shield_rect,
+                    asset_map["shield"].semantic_color,
+                    asset_map["shield"].asset,
+                    gen_original=gen_original,
+                    ignore_mask=bbox_smap_for_agent,
+                    cut_mask_top_ratio=0.45,
+                )
+
+    return img

multimodal/examples/mugen/data/coinrun/game.py

@@ -0,0 +1,295 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+
+
+class Game:
+    def __init__(self, **kwargs):
+        self.game_id = -1
+        self.level_seed = 0
+        self.rl_agent_seed = 0
+        self.zoom = 5.5
+        self.bgzoom = 0.4  # NOTE: hard-coded
+        self.world_theme_n = -1
+        self.agent_theme_n = -1
+
+        self.background_themes = []
+        self.ground_themes = []
+        self.agent_themes = []
+        self.monster_names = {}
+        self.flattened_monster_names = []
+
+        # TODO: save and load these from the game engine
+        self.video_res = 1024
+        self.maze_w = 64
+        self.maze_h = 13  # for zoom 5.5
+
+        self.reset_game()
+
+        self.__dict__.update(**kwargs)
+        self.frames = [Frame(**f) for f in self.frames]
+
+    def reset_game(self):
+        self.maze = None
+        self.frames = []
+
+    def asdict(self, f_start=-1, f_end=-1):
+        if f_end < 0:
+            # show all frames by default
+            frames_as_dict = [f.asdict() for f in self.frames]
+        else:
+            frames_as_dict = [f.asdict() for f in self.frames[f_start:f_end]]
+        return {
+            "game_id": self.game_id,
+            "level_seed": self.level_seed,
+            "rl_agent_seed": self.rl_agent_seed,
+            "zoom": self.zoom,
+            "bgzoom": self.bgzoom,
+            "world_theme_n": self.world_theme_n,
+            "agent_theme_n": self.agent_theme_n,
+            "background_themes": self.background_themes,
+            "ground_themes": self.ground_themes,
+            "agent_themes": self.agent_themes,
+            "monster_names": self.monster_names,
+            "video_res": self.video_res,
+            "maze_w": self.maze_w,
+            "maze_h": self.maze_h,
+            "maze": self.maze if self.maze is not None else None,
+            "frames": frames_as_dict,
+        }
+
+    def __repr__(self):
+        return json.dumps(self.asdict())
+
+    def save_json(self, json_path, f_start=-1, f_end=-1):
+        with open(json_path, "w") as f:
+            json.dump(self.asdict(f_start, f_end), f, indent=2)
+
+    def load_json(self, json_path):
+        with open(json_path, "r") as f:
+            data = json.load(f)
+
+        self.reset_game()
+        self.__dict__.update(**data)
+        self.frames = [Frame(**f) for f in self.frames]
+
+        self.flatten_monster_names()
+        self.reset_eaten_coins()
+
+    def flatten_monster_names(self):
+        # the order is important!
+        self.flattened_monster_names = self.monster_names["ground"]
+        self.flattened_monster_names.extend(self.monster_names["walking"])
+        self.flattened_monster_names.extend(self.monster_names["flying"])
+
+    # NOTE: some coins might be missing due to how 3s clip json is saved
+    # reset all eaten coins to put them back
+    # this is a temporary fix until we regenerate all jsons
+    def reset_eaten_coins(self):
+        for coin_loc in self.frames[-1].coins_eaten:
+            # note the game rows are saved as strings
+            # NOTE: '1' is the yellow coin, we also has another type '2' that is the red gem
+            # but the json with '2' enabled should not have this issue
+            if self.maze[coin_loc[1]][coin_loc[0]] == ".":
+                self.maze[coin_loc[1]] = (
+                    self.maze[coin_loc[1]][: coin_loc[0]]
+                    + "1"
+                    + self.maze[coin_loc[1]][(coin_loc[0] + 1) :]
+                )
+
+
+class Frame:
+    def __init__(self, **kwargs):
+        self.frame_id = -1
+        self.file_name = ""
+        self.state_time = 0
+        self.coins_eaten = []
+        self.agent = None
+        self.monsters = []
+
+        self.__dict__.update(**kwargs)
+        if "agent" in self.__dict__ and self.agent is not None:
+            self.agent = Agent(**self.agent)
+        if "monsters" in self.__dict__:
+            self.monsters = [Monster(**m) for m in self.monsters]
+
+    def asdict(self):
+        return {
+            "frame_id": self.frame_id,
+            "file_name": self.file_name,
+            "state_time": self.state_time,
+            "coins_eaten": self.coins_eaten,
+            "agent": self.agent.asdict() if self.agent is not None else None,
+            "monsters": [m.asdict() for m in self.monsters],
+        }
+
+    def __repr__(self):
+        return json.dumps(self.asdict())
+
+
+class Agent:
+    def __init__(
+        self,
+        x,
+        y,
+        vx=0.0,
+        vy=0.0,
+        time_alive=0,
+        ladder=False,
+        spring=0,
+        is_killed=False,
+        killed_animation_frame_cnt=0,
+        finished_level_frame_cnt=0,
+        killed_monster=False,
+        bumped_head=False,
+        collected_coin=False,
+        collected_gem=False,
+        power_up_mode=False,
+        **kwargs,
+    ):
+        self.x = x
+        self.y = y
+        self.vx = vx
+        self.vy = vy
+        self.time_alive = time_alive
+        self.ladder = ladder  # for climb pose
+        self.spring = spring  # for duck pose
+
+        # states related to agent dying or finishing animations
+        self.is_killed = is_killed
+        self.killed_animation_frame_cnt = killed_animation_frame_cnt
+        self.finished_level_frame_cnt = finished_level_frame_cnt
+        self.killed_monster = killed_monster
+        self.bumped_head = bumped_head
+        self.collected_coin = collected_coin
+        self.collected_gem = collected_gem
+        self.power_up_mode = power_up_mode
+
+        self.anim_freq = 5  # hard-coded
+
+        # decide whether to flip asset horizontally
+        self.is_facing_right = True
+        if self.vx < 0:
+            self.is_facing_right = False
+
+        # decide which of the two walk/climb asset to use
+        self.walk1_mode = True
+        if (self.time_alive // self.anim_freq) % 2 != 0:
+            self.walk1_mode = False
+
+        self.pose = self.get_pose()
+
+        # kwargs are ignored
+        # self.__dict__.update(**kwargs)
+
+    def get_pose(self):
+        if self.is_killed:
+            return "hit"
+        if self.ladder:
+            if self.walk1_mode:
+                return "climb1"
+            else:
+                return "climb2"
+        if self.vy != 0:
+            return "jump"
+        if self.spring != 0:
+            return "duck"
+        if self.vx == 0:
+            return "stand"
+        if self.walk1_mode:
+            return "walk1"
+        else:
+            return "walk2"
+
+    def asdict(self):
+        return {
+            "x": self.x,
+            "y": self.y,
+            "vx": self.vx,
+            "vy": self.vy,
+            "time_alive": self.time_alive,
+            "ladder": self.ladder,
+            "spring": self.spring,
+            "is_killed": self.is_killed,
+            "killed_animation_frame_cnt": self.killed_animation_frame_cnt,
+            "finished_level_frame_cnt": self.finished_level_frame_cnt,
+            "killed_monster": self.killed_monster,
+            "bumped_head": self.bumped_head,
+            "collected_coin": self.collected_coin,
+            "collected_gem": self.collected_gem,
+            "power_up_mode": self.power_up_mode,
+            "anim_freq": self.anim_freq,
+            "is_facing_right": self.is_facing_right,
+            "walk1_mode": self.walk1_mode,
+            "pose": self.pose,
+        }
+
+    def __repr__(self):
+        return json.dumps(self.asdict())
+
+
+class Monster:
+    def __init__(
+        self,
+        m_id,
+        x,
+        y,
+        vx=0.0,
+        vy=0.0,
+        theme=0,
+        is_flying=False,
+        is_walking=False,
+        is_jumping=False,
+        is_dead=False,
+        time=0,
+        anim_freq=1,
+        monster_dying_frame_cnt=0,
+        **kwargs,
+    ):
+        self.m_id = m_id
+        self.x = x
+        self.y = y
+        self.vx = vx
+        self.vy = vy
+        self.theme = theme  # monster type (saw, snail, slime, etc.)
+        self.is_flying = is_flying
+        self.is_walking = is_walking
+        self.is_jumping = is_jumping
+        self.is_dead = is_dead
+        self.time = time
+        self.anim_freq = anim_freq
+        self.monster_dying_frame_cnt = monster_dying_frame_cnt
+
+        # decide which of the two walk/climb asset to use
+        self.walk1_mode = True
+        if self.is_jumping:
+            # for jumping monster, walk1 asset is decided by vertical speed
+            if self.vy != 0:
+                self.walk1_mode = False
+        elif (self.time // self.anim_freq) % 2 != 0:
+            self.walk1_mode = False
+
+    def asdict(self):
+        return {
+            "m_id": self.m_id,
+            "x": self.x,
+            "y": self.y,
+            "vx": self.vx,
+            "vy": self.vy,
+            "theme": self.theme,
+            "is_flying": self.is_flying,
+            "is_walking": self.is_walking,
+            "is_jumping": self.is_jumping,
+            "is_dead": self.is_dead,
+            "time": self.time,
+            "anim_freq": self.anim_freq,
+            "monster_dying_frame_cnt": self.monster_dying_frame_cnt,
+            "walk1_mode": self.walk1_mode,
+        }
+
+    def __repr__(self):
+        return json.dumps(self.asdict())

multimodal/examples/mugen/data/coinrun/generate_text_desc.py

@@ -0,0 +1,435 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+
+
+class Sequence:
+    def __init__(
+        self, start_frame, end_frame, pose_type, start_x, start_y, end_x, end_y
+    ):
+        self.start_frame = start_frame
+        self.end_frame = end_frame
+
+        # 'ground' includes 'walk', 'duck', 'stand'; other types are 'climb', 'jump', 'hit'
+        self.pose_type = pose_type
+        self.start_x = start_x
+        self.start_y = start_y
+        self.end_x = end_x
+        self.end_y = end_y
+        self.time_jumps = 1 if pose_type == "jump" else 0
+        self.end_maze_above = "."
+        self.end_maze_below = "."
+        self.num_coins_eaten = 0
+        self.num_gems_eaten = 0
+        self.start_shield = False
+        self.end_shield = False
+        self.changed_shield = False
+        self.killed_monsters = []
+        self.jump_over_monsters = []
+        self.killed_by = ""
+        self.text_desc = ""
+
+        # Decide graduarity of text description (skip sequence shorter than this)
+        self.min_len_for_text_desc = 5
+
+    def asdict(self):
+        return {
+            "start_frame": self.start_frame,
+            "end_frame": self.end_frame,
+            "pose_type": self.pose_type,
+            "start_xy": (self.start_x, self.start_y),
+            "end_xy": (self.end_x, self.end_y),
+            "bumped_head": self.is_bumped_head(),
+            "same_level_jump": self.is_same_level_jump(),
+            "num_coins_eaten": self.num_coins_eaten,
+            "num_gems_eaten": self.num_gems_eaten,
+            "start_shield": self.start_shield,
+            "end_shield": self.end_shield,
+            "changed_shield": self.changed_shield,
+            "killed_monsters": self.killed_monsters,
+            "jump_over_monsters": self.jump_over_monsters,
+            "killed_by": self.killed_by,
+            "text_desc": self.text_desc,
+        }
+
+    def __repr__(self):
+        return json.dumps(self.asdict())
+
+    # bumped head will show as 'walk' pose and last for 1-2 frames
+    def is_bumped_head(self):
+        if (
+            self.pose_type == "ground"
+            and (self.end_frame - self.start_frame <= 1)
+            and self.end_maze_below in ".12"
+        ):  # and self.end_maze_above in 'SAab'
+            return True
+        return False
+
+    def is_same_level_jump(self):
+        if self.pose_type == "jump" and abs(self.end_y - self.start_y) <= 0.5:
+            return True
+        return False
+
+    def merge_sequences(self, sequences):
+        self.end_frame = sequences[-1].end_frame
+        self.end_x = sequences[-1].end_x
+        self.end_y = sequences[-1].end_y
+        self.end_maze_above = sequences[-1].end_maze_above
+        self.end_maze_below = sequences[-1].end_maze_below
+        for seq in sequences:
+            if seq.is_bumped_head():
+                self.time_jumps -= 1
+            self.time_jumps += seq.time_jumps
+
+            self.num_coins_eaten += seq.num_coins_eaten
+            self.num_gems_eaten += seq.num_gems_eaten
+            self.killed_monsters.extend(seq.killed_monsters)
+            self.jump_over_monsters.extend(seq.jump_over_monsters)
+
+    def process_metadata(self, game):
+        # generate game.flattened_monster_names if not already
+        # this is used to get monster names
+        if len(game.flattened_monster_names) == 0:
+            game.flatten_monster_names()
+
+        # count number of coins and gems eaten during the sequence
+        # start from one frame earlier (if not 0) so we can get change in the first frame
+        start_frame_id = max(self.start_frame - 1, 0)
+        if len(game.frames[self.end_frame].coins_eaten) > len(
+            game.frames[start_frame_id].coins_eaten
+        ):
+            start_coin_set = {
+                (coord[0], coord[1])
+                for coord in game.frames[start_frame_id].coins_eaten
+            }
+            end_coin_set = {
+                (coord[0], coord[1])
+                for coord in game.frames[self.end_frame].coins_eaten
+            }
+            new_coins_eaten = end_coin_set - start_coin_set
+            for coin_coord in new_coins_eaten:
+                if game.maze[coin_coord[1]][coin_coord[0]] == "2":
+                    self.num_gems_eaten += 1
+                else:
+                    self.num_coins_eaten += 1
+
+        # check if Mugen changes between shield up and down mode during the sequence
+        self.start_shield = game.frames[self.start_frame].agent.power_up_mode
+        self.end_shield = game.frames[self.end_frame].agent.power_up_mode
+        shield_up_mode = False
+        shield_down_mode = False
+        for frame_id in range(self.start_frame, self.end_frame + 1):
+            if game.frames[frame_id].agent.power_up_mode:
+                shield_up_mode = True
+            else:
+                shield_down_mode = True
+        if shield_up_mode and shield_down_mode:
+            self.changed_shield = True
+
+        end_frame_id = min(self.end_frame + 2, len(game.frames))
+        for frame_id in range(self.start_frame, end_frame_id):
+            frame = game.frames[frame_id]
+            dead_monsters = set()
+            for i, m in enumerate(frame.monsters):
+                if m.is_dead:
+                    dead_monsters.add(i)
+            # if more monsters are killed, record the monster killed and the frame id
+            if frame_id > self.start_frame and len(dead_monsters) > len(
+                prev_dead_monsters
+            ):
+                killed_monster_theme = frame.monsters[
+                    list(dead_monsters - prev_dead_monsters)[0]
+                ].theme
+                self.killed_monsters.append(
+                    game.flattened_monster_names[killed_monster_theme]
+                )
+            prev_dead_monsters = dead_monsters.copy()
+
+        # figure out which monster killed Mugen
+        killed_by_m_id = -1
+        if self.pose_type == "hit":
+            # check the monster distance in the first frame of hit sequence
+            m_min_dist = 1000  # just put some random large dist here
+            for m in game.frames[self.start_frame].monsters:
+                x_dist = self.start_x - m.x
+                y_dist = self.start_y - m.y
+                m_dist = x_dist * x_dist + y_dist * y_dist
+                if m_dist < m_min_dist:
+                    killed_by_m_id = m.theme
+                    m_min_dist = m_dist
+            if killed_by_m_id != -1:
+                self.killed_by = game.flattened_monster_names[killed_by_m_id]
+
+        # check for monsters jumped over
+        if self.pose_type == "jump":
+            # for purpose of checking jumped over monsters,
+            # ground y is fixed at the y coordinate of the previous frame
+            # note for jump sequence, start_y already recorded the location before jump starts
+            ground_y = round(self.start_y)
+            jump_over_monsters_set = set()
+            for frame_id in range(self.start_frame, self.end_frame + 1):
+                frame = game.frames[frame_id]
+                # this is the location below the agent at the same y level when jump starts
+                ground_loc = (round(frame.agent.x), ground_y)
+                for i, m in enumerate(frame.monsters):
+                    if (round(m.x), round(m.y)) == ground_loc:
+                        # use set to avoid adding duplicates
+                        jump_over_monsters_set.add(i)
+
+            # now convert these into names, but only keep those that's still not killed by the next frame
+            for m_i in jump_over_monsters_set:
+                if not game.frames[end_frame_id - 1].monsters[m_i].is_dead:
+                    self.jump_over_monsters.append(
+                        game.flattened_monster_names[frame.monsters[m_i].theme]
+                    )
+
+    def generate_text_desc(self):
+        # only generate if sequence is long enough
+        if self.end_frame - self.start_frame < self.min_len_for_text_desc:
+            self.text_desc = ""
+        elif self.pose_type == "hit":
+            if self.killed_by != "":
+                self.text_desc = f"killed by a {self.killed_by}"
+            else:
+                self.text_desc = "killed by a monster"
+        else:
+            y_direct = ""
+            if self.end_y - self.start_y > 0.5:
+                y_direct = " up"
+            elif self.start_y - self.end_y > 0.5:
+                y_direct = " down"
+            else:
+                y_direct = " a bit" if self.pose_type == "ground" else ""
+            x_direct = ""
+            if self.end_x - self.start_x > 0.5:
+                x_direct = " to the right"
+            elif self.start_x - self.end_x > 0.5:
+                x_direct = " to the left"
+            else:
+                x_direct = " a bit" if self.pose_type == "ground" else ""
+
+            if self.pose_type == "climb":
+                self.text_desc = f"climbs{y_direct} on a ladder"
+            elif self.pose_type == "ground":
+                self.text_desc = f"walks{x_direct}"  # TODO: add random verbs
+            elif self.pose_type == "jump":
+                jump_time_desc = ""
+                if self.time_jumps >= 2:
+                    jump_time_desc = " a few times"
+
+                # only add jump destination if it's not a same level jump
+                jump_dest_desc = ""
+                if y_direct != "":
+                    if self.end_maze_below in "SAab":
+                        if self.end_y < 1.5:
+                            jump_dest_desc = " to the ground"
+                        else:
+                            jump_dest_desc = " to a platform"
+                    elif self.end_maze_below in "#$&%":
+                        jump_dest_desc = " to a crate"
+                    elif self.end_maze_below == "=":
+                        jump_dest_desc = " to a ladder"
+
+                # add desc for monsters jumped over
+                jumped_over_desc = ""
+                if len(self.jump_over_monsters) > 0:
+                    jumped_over_desc = " over a " + " and a ".join(
+                        self.jump_over_monsters
+                    )
+
+                self.text_desc = f"jumps{y_direct}{jump_time_desc}{x_direct}{jumped_over_desc}{jump_dest_desc}"
+
+            if self.num_coins_eaten > 0 or self.num_gems_eaten > 0:
+                self.text_desc += self.generate_collect_coin_desc()
+
+            if len(self.killed_monsters) > 0:
+                self.text_desc += " and killed a " + " and a ".join(
+                    self.killed_monsters
+                )
+
+    def generate_collect_coin_desc(self):
+        if self.num_coins_eaten == 0 and self.num_gems_eaten == 0:
+            return ""
+
+        coin_descs = []
+        # add coin description if collected at least one coin
+        if self.num_coins_eaten == 1:
+            coin_descs.append(" a coin")
+        elif self.num_coins_eaten > 1:
+            coin_descs.append(" a few coins")
+
+        # add gem description if collected at least one gem
+        if self.num_gems_eaten == 1:
+            coin_descs.append(" a gem")
+        elif self.num_gems_eaten > 1:
+            coin_descs.append(" a few gems")
+
+        # connects descriptions for coins and gems with 'and'
+        coin_descs = " and".join(coin_descs)
+
+        # shield change should only be a result of eating gem or coin
+        if self.changed_shield:
+            coin_descs += self.generate_shield_desc()
+
+        return f" and collects{coin_descs}"
+
+    def generate_shield_desc(self):
+        if not self.start_shield and self.end_shield:
+            return " to turn on the shield"
+        elif self.start_shield and not self.end_shield:
+            return " to turn off the shield"
+        else:
+            # start and end in the same shield state but still changed shield during sequence
+            if self.start_shield:
+                return " to turn shield off then on again"
+            else:
+                return " to turn shield on then off again"
+
+
+def process_sequence(game, curr_pose_type, start_i, curr_i, last_seq=False):
+    # different type of pose, construct a sequence
+    # for 'jump', the start and end location is based on frame before the first and after the last frame
+    # for others, it's the first and last frame
+    if curr_pose_type == "jump":
+        pos_start_frame = max(start_i - 1, 0)
+        pos_end_frame = curr_i
+    else:
+        pos_start_frame = start_i
+        # curr_i will be one frame after, unless it's the last sequence of video
+        # however, for jump sequence, we do want one frame after to know where jump lands
+        pos_end_frame = curr_i - 1 if not last_seq else curr_i
+
+    seq = Sequence(
+        start_frame=start_i,
+        end_frame=curr_i - 1 if not last_seq else curr_i,
+        pose_type=curr_pose_type,
+        start_x=game.frames[pos_start_frame].agent.x,
+        start_y=game.frames[pos_start_frame].agent.y,
+        end_x=game.frames[pos_end_frame].agent.x,
+        end_y=game.frames[pos_end_frame].agent.y,
+    )
+    seq.end_maze_above = game.maze[round(seq.end_y) + 1][round(seq.end_x)]
+    seq.end_maze_below = game.maze[round(seq.end_y) - 1][round(seq.end_x)]
+    # sometimes jump may end a bit over the edge of cliff, this is to catch and fix that
+    if curr_pose_type == "jump" and seq.end_maze_below in ".12":
+        neighbor_x = (
+            int(seq.end_x) * 2 + 1 - round(seq.end_x)
+        )  # get the opposite of round()
+        seq.end_maze_below = game.maze[round(seq.end_y) - 1][neighbor_x]
+
+    return seq
+
+
+def convert_game_to_text_desc(game, start_idx=0, end_idx=-1, alien_name="Mugen"):
+    if alien_name is None:
+        alien_name = "Mugen"
+
+    # if end_idx is not specified, set it to end of the game level
+    if end_idx == -1:
+        end_idx = len(game.frames)
+    start_idx = max(0, start_idx)
+    end_idx = min(len(game.frames), end_idx)
+
+    sequences = []
+    for i, f in enumerate(game.frames[start_idx:end_idx]):
+        pose = f.agent.pose.strip("12")
+        if pose in ["walk", "duck", "stand"]:
+            pose_type = "ground"
+        else:
+            pose_type = pose
+        if i == 0:
+            # first frame, initialize some status
+            start_i = 0
+            curr_pose_type = pose_type
+            continue
+
+        if pose_type == curr_pose_type:
+            # same type of pose, same sequence
+            continue
+        else:
+            seq = process_sequence(
+                game, curr_pose_type, start_idx + start_i, start_idx + i, last_seq=False
+            )
+            sequences.append(seq)
+            start_i = i
+            curr_pose_type = pose_type
+
+    # add the last leftover sequence
+    seq = process_sequence(
+        game, curr_pose_type, start_idx + start_i, start_idx + i, last_seq=True
+    )
+    sequences.append(seq)
+
+    # collapse two jumps into one sequence
+    # first pass, merge jumps before and after bumped head, this is to correctly identify jumps at the same level
+    seq_i = 0
+    reduced_sequences = []
+    while seq_i < len(sequences):
+        if seq_i == 0 or seq_i == len(sequences) - 1:
+            reduced_sequences.append(sequences[seq_i])
+            seq_i += 1
+        elif (
+            sequences[seq_i].is_bumped_head()
+            and reduced_sequences[-1].pose_type == "jump"
+            and sequences[seq_i + 1].pose_type == "jump"
+        ):
+            # in case of bumped head, merge the jumps before and after
+            reduced_sequences[-1].merge_sequences(sequences[seq_i : seq_i + 2])
+            seq_i += 2
+        else:
+            reduced_sequences.append(sequences[seq_i])
+            seq_i += 1
+    sequences = reduced_sequences
+
+    # second pass, collapse two jumps into one sequence if they're both same level jumps
+    # jump up and down are not merged (unless it's separated by bumped head that will be merged in first pass)
+    result_sequences = []
+    seq_i = 0
+    max_ground_seq_len_to_merge = 5
+    while seq_i < len(sequences):
+        # only merge if it's a 'ground' sequence, and before/after are both jumps
+        if (
+            sequences[seq_i].pose_type != "ground"
+            or seq_i == 0
+            or seq_i == len(sequences) - 1
+        ):
+            result_sequences.append(sequences[seq_i])
+            seq_i += 1
+        elif (
+            result_sequences[-1].pose_type != "jump"
+            or sequences[seq_i + 1].pose_type != "jump"
+        ):
+            result_sequences.append(sequences[seq_i])
+            seq_i += 1
+        elif (
+            result_sequences[-1].is_same_level_jump()
+            and sequences[seq_i + 1].is_same_level_jump()
+            and (
+                sequences[seq_i].end_frame - sequences[seq_i].start_frame
+                < max_ground_seq_len_to_merge
+            )
+        ):
+            # not bumped head, then only merge if sequence is short enough, and both jumps are the same level
+            result_sequences[-1].merge_sequences(sequences[seq_i : seq_i + 2])
+            seq_i += 2
+        else:
+            result_sequences.append(sequences[seq_i])
+            seq_i += 1
+    sequences = result_sequences
+
+    # generate text description for each sequence
+    text_descriptions = []
+    for seq in sequences:
+        seq.process_metadata(game)
+        seq.generate_text_desc()
+        if seq.text_desc != "":
+            text_descriptions.append(seq.text_desc)
+
+    # add Mugen in the beginning, then concat by 'and'
+    final_text_desc = alien_name + " " + ", and ".join(text_descriptions)
+
+    return final_text_desc

multimodal/examples/mugen/data/mugen_datamodules.py

@@ -0,0 +1,112 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional
+
+import pytorch_lightning as pl
+import torch
+import torch.distributed as dist
+import torch.utils.data as data
+
+from .mugen_dataset import MUGENDataset, MUGENDatasetArgs
+
+
+class MUGENDataModule(pl.LightningDataModule):
+    """General lightning data module for MUGEN dataset.
+
+    Args:
+        mugen_dataset_args (MUGENDatasetArgs): arguments for MUGENDataset.
+        text_transform (Optional[Callable]): transform for text batches.
+            Only used when not ``None`` and when ``mugen_dataset_args.get_text_desc = True``.
+            Defaults to ``None``.
+        video_transform (Optional[Callable]): transform for video batches.
+            Only used when not ``None`` and when ``mugen_dataset_args.get_game_frame = True``.
+            Defaults to ``None``.
+        audio_transform (Optional[Callable]): transform for audio batches.
+            Only used when not ``None`` and when ``mugen_dataset_args.get_audio = True``.
+            Defaults to ``None``.
+        batch_size (int): number of samples per batch.
+            Defaults to ``16``.
+        num_workers (int): number of subprocesses for data loading.
+            Defaults to ``0``, meaning data is loaded in the main process.
+        shuffle (bool): whether to reshuffle data after each epoch.
+            Defaults to ``True``.
+    """
+
+    def __init__(
+        self,
+        mugen_dataset_args: MUGENDatasetArgs,
+        text_transform: Optional[Callable] = None,
+        video_transform: Optional[Callable] = None,
+        audio_transform: Optional[Callable] = None,
+        batch_size: int = 16,
+        num_workers: int = 0,
+        shuffle: bool = True,
+    ):
+        super().__init__()
+        self.mugen_dataset_args = mugen_dataset_args
+        self.text_transform = text_transform
+        self.video_transform = video_transform
+        self.audio_transform = audio_transform
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.shuffle = shuffle
+
+    @property
+    def n_classes(self):
+        dataset = self._dataset(True)
+        return dataset.n_classes
+
+    def _custom_collate_fn(self, batch):
+        collated_batch = {}
+        if self.mugen_dataset_args.get_game_frame:
+            video = [elem["video"] for elem in batch]
+            video = torch.stack(video)
+            video = self.video_transform(video) if self.video_transform else video
+            collated_batch["video"] = video
+        if self.mugen_dataset_args.get_text_desc:
+            text = [elem["text"] for elem in batch]
+            # cannot be torch.stack'ed because still in raw text form, not Tensor
+            text = self.text_transform(text) if self.text_transform else text
+            collated_batch["text"] = text
+        if self.mugen_dataset_args.get_audio:
+            audio = [elem["audio"] for elem in batch]
+            audio = torch.stack(audio)
+            audio = self.audio_transform(audio) if self.audio_transform else audio
+            collated_batch["audio"] = audio
+        return collated_batch
+
+    def _dataset(self, split):
+        dataset = MUGENDataset(args=self.mugen_dataset_args, split=split)
+        return dataset
+
+    def _dataloader(self, split):
+        dataset = self._dataset(split)
+        if dist.is_initialized():
+            sampler = data.distributed.DistributedSampler(
+                dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank()
+            )
+        else:
+            sampler = None
+        dataloader = data.DataLoader(
+            dataset,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            sampler=sampler,
+            shuffle=sampler is None and self.shuffle is True,
+            collate_fn=self._custom_collate_fn,
+        )
+        return dataloader
+
+    def train_dataloader(self):
+        return self._dataloader("train")
+
+    def val_dataloader(self):
+        return self._dataloader("val")
+
+    def test_dataloader(self):
+        return self._dataloader("test")

multimodal/examples/mugen/generation/LoadAndComparePretrainedVQVAE.ipynb

@@ -0,0 +1,383 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ee3d68e4",
+   "metadata": {},
+   "source": [
+    "# Compare MUGEN's Video VQVAE with TorchMultimodal's\n",
+    "\n",
+    "This notebook loads the public MUGEN checkpoint for Video VQVAE, remaps the state_dict, and loads it into TorchMultimodal's Video VQVAE to ensure the outputs match. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5af9d001",
+   "metadata": {},
+   "source": [
+    "### Set directories\n",
+    "\n",
+    "Replace these with your local directories."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "071c8b48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "checkpoint_dir = '/Users/rafiayub/checkpoints/'\n",
+    "repo_dir = '/Users/rafiayub/mugen/'\n",
+    "home_dir = '/Users/rafiayub/'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a3a0f19f",
+   "metadata": {},
+   "source": [
+    "### Clone MUGEN's repo"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "83812502",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!git clone https://github.com/mugen-org/MUGEN_baseline.git $repo_dir"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "07757cfa",
+   "metadata": {},
+   "source": [
+    "### Download and unzip checkpoints\n",
+    "\n",
+    "This will take some time."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d41a0c86",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!wget https://dl.noahmt.com/creativity/data/MUGEN_release/checkpoints.zip -P $checkpoint_dir"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01d9638a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Unzip checkpoints\n",
+    "zip_location = os.path.join(checkpoint_dir, 'checkpoints.zip')\n",
+    "!unzip $zip_location -d $checkpoint_dir"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f06c8938",
+   "metadata": {},
+   "source": [
+    "### Load checkpoint into MUGEN model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f3e74b3a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "sys.path.append(home_dir)\n",
+    "\n",
+    "import torch\n",
+    "from torch import nn\n",
+    "import mugen\n",
+    "\n",
+    "ckpt = torch.load(\n",
+    "    os.path.join(checkpoint_dir, 'generation/video_vqvae/L32/epoch=54-step=599999.ckpt'), \n",
+    "    map_location=torch.device('cpu')\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3ea6d13e",
+   "metadata": {},
+   "source": [
+    "The arguments are taken from MUGEN's training scripts found at: https://github.com/mugen-org/MUGEN_baseline/blob/main/generation/experiments/vqvae/VideoVQVAE_L32.sh"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "f81bea2e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class Namespace:\n",
+    "    def __init__(self, **kwargs):\n",
+    "        self.__dict__.update(kwargs)\n",
+    "\n",
+    "\n",
+    "vqvae_args=Namespace(\n",
+    "    embedding_dim=256,\n",
+    "    n_codes=2048,\n",
+    "    n_hiddens=240,\n",
+    "    n_res_layers=4,\n",
+    "    lr=0.0003,\n",
+    "    downsample=(4, 32, 32),\n",
+    "    kernel_size=3,\n",
+    "    sequence_length=16,\n",
+    "    resolution=256,\n",
+    ")\n",
+    "vv_mugen = mugen.VQVAE(vqvae_args)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "fbdcf1f6",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "vv_mugen.load_state_dict(ckpt['state_dict'])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a6bfb325",
+   "metadata": {},
+   "source": [
+    "### Create TorchMultimodal's Video VQVAE"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "74e6bd54",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from examples.mugen.generation.video_vqvae import video_vqvae_mugen\n",
+    "\n",
+    "vv_torchmm = video_vqvae_mugen(pretrained_model_key=None)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e612d831",
+   "metadata": {},
+   "source": [
+    "### Remap MUGEN's state_dict and load into new model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "5f4d4774",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import re\n",
+    "\n",
+    "def map_state_dict(state_dict):\n",
+    "    mapped_state_dict = {}\n",
+    "    dim_map = {'w': '2', 'h': '1', 't': '0'}\n",
+    "    layer_map = {'w_qs': 'query', 'w_ks': 'key', 'w_vs': 'value', 'fc': 'output'}\n",
+    "    for param, val in state_dict.items():\n",
+    "        new_param = param\n",
+    "        res = re.search('encoder.convs.', param)\n",
+    "        if res:\n",
+    "            idx = res.end()\n",
+    "            layer_id = int(param[idx])\n",
+    "            new_param = param[:idx] + str(layer_id * 2) + param[idx+1:]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        res = re.search('encoder.conv_last', param)\n",
+    "        if res:\n",
+    "            idx = res.start() + len('encoder.')\n",
+    "            new_param = param[:idx] + 'convs.10' + param[res.end():]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        res = re.search('attn_[w,h,t]\\..*\\.', param)\n",
+    "        if res:\n",
+    "            dim = param[res.start()+5]\n",
+    "            new_dim = dim_map[dim]\n",
+    "            layer = param[res.start()+7:res.end()-1]\n",
+    "            new_layer = layer_map[layer]\n",
+    "            new_param = param[:res.start()] + 'mha_attns.' + new_dim + '.' + new_layer + '.' + param[res.end():]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        res = re.search('pre_vq_conv', param)\n",
+    "        if res:\n",
+    "            new_param = 'encoder.conv_out' + param[res.end():]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        res = re.search('post_vq_conv', param)\n",
+    "        if res:\n",
+    "            new_param = 'decoder.conv_in' + param[res.end():]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        res = re.search('decoder.convts.', param)\n",
+    "        if res:\n",
+    "            idx = res.end()\n",
+    "            layer_id = int(param[idx])\n",
+    "            new_param = param[:idx] + str(layer_id * 2) + param[idx+1:]\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        if param == 'codebook.N':\n",
+    "            new_param = 'codebook.code_usage'\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        if param == 'codebook.z_avg':\n",
+    "            new_param = 'codebook.code_avg'\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "        if param == 'codebook.embeddings':\n",
+    "            new_param = 'codebook.embedding'\n",
+    "            mapped_state_dict[new_param] = val\n",
+    "            continue\n",
+    "            \n",
+    "        mapped_state_dict[new_param] = val\n",
+    "        \n",
+    "    return mapped_state_dict"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "38234858",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "new_state_dict = map_state_dict(ckpt['state_dict'])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "e160fb51",
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "vv_torchmm.load_state_dict(new_state_dict)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "46d58eb7",
+   "metadata": {},
+   "source": [
+    "### Compare outputs with a random input"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "3c85cdd3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Max difference between outputs: 3.0875205993652344e-05\n",
+      "Mean difference between outputs: 1.7353995929170196e-07\n"
+     ]
+    }
+   ],
+   "source": [
+    "torch.manual_seed(4)\n",
+    "video = torch.randn(1,3,32,256,256) # b, c, t, h, w\n",
+    "\n",
+    "vv_mugen.eval()\n",
+    "vv_torchmm.eval()\n",
+    "\n",
+    "loss, x_recon, codebook_output = vv_mugen(video)\n",
+    "output = vv_torchmm(video)\n",
+    "\n",
+    "diff = abs(output.decoded - x_recon)\n",
+    "print(f'Max difference between outputs: {torch.max(diff).item()}')\n",
+    "print(f'Mean difference between outputs: {torch.mean(diff).item()}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fa78569e",
+   "metadata": {},
+   "source": [
+    "### Save mapped checkpoint"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "48651d44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "save_path = '/Users/rafiayub/checkpoints/generation/video_vqvae/mugen_video_vqvae_L32.pt'\n",
+    "torch.save(new_state_dict, save_path)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

multimodal/examples/mugen/generation/README.md

@@ -0,0 +1,33 @@
+# Text-to-Video Generation with MUGEN
+
+This directory contains the high-level model components for text-to-video generation following [MUGEN](https://arxiv.org/abs/2204.08058). They demonstrate how to use building blocks from TorchMultimodal to quickly assemble a new auto-regressive generative model for different pairs of modalities. Here is a [colab demo](https://colab.research.google.com/drive/1C3ZbH_l19g_KqW3CPeX2-8Q2sOUCpmZo?usp=sharing) showing how to generate a video clip from text prompts.
+
+https://user-images.githubusercontent.com/23155714/196074330-6f03593c-da8e-473f-8935-8bf1950baa33.mp4
+
+```python
+from torchmultimodal.utils.generate import GenerationUtil
+from examples.mugen.generation.text_video_gpt import text_video_gpt
+
+
+model = text_video_gpt(video_seq_len=32, pretrained_text_video_gpt_model_key="mugen_L32")
+generator = GenerationUtil(model)
+
+output = generator.sample(
+    ['Mugen moves left to right on a cliff and picks up a gem.'],
+    max_seq_len=512,
+    use_cache=True,
+    causal=True,
+    device=<current_device>,
+)
+samples = output.decoded
+```
+
+## Model
+The model architecture used by MUGEN follows [DALL-E](https://arxiv.org/abs/2102.12092) but with the image components replaced by those for video following [VideoGPT](https://arxiv.org/abs/2104.10157).
+
+Multimodal generation involves generation of samples in one modality given inputs from another. As in the text-to-image generation model DALL-E, it typically involves a two-stage process of first learning a discrete latent representation for each modality and then using a [GPT](https://openai.com/blog/language-unsupervised/) transformer decoder to learn a joint prior for both modalities in the latent space. For text data, the latent representation is obtained through tokenization such as [BPE](https://en.wikipedia.org/wiki/Byte_pair_encoding) used in this example. For high dimensional data such as video and image, a [VQ-VAE](https://arxiv.org/abs/1711.00937) model is used to learn a set of downsampled discrete embedding vectors through nearest-neighbor lookups from a "codebook" where the chosen indices are referred to as the token ids following convention from language modeling.
+
+VideoGPT is a generative model for video using a VQ-VAE model with video encoder/decoder and a GPT transformer decoder for token generation. The encoder and the decoder use 3D-convolution and self axial-attention to learn video information.
+
+## Generation
+In this example generation refers to the auto-regressive process where we iteratively predict the next token id from the current until reaching the desired output length, a technique initially used by language modeling but has been extended to multimodal generation. To control the generation process, a top level abstraction is provided as a utility in [generate.py](https://github.com/facebookresearch/multimodal/blob/main/torchmultimodal/utils/generate.py) which takes the model as an input.

multimodal/examples/mugen/generation/text_video_gpt.py

@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List, Optional, Tuple
+
+import torch
+
+from examples.mugen.generation.video_vqvae import video_vqvae_mugen
+
+from torch import nn, Tensor
+
+from torchmultimodal.models.video_gpt.gpt import (
+    MultimodalGPT,
+    MultimodalTransformerDecoder,
+    RightShift,
+    TransformerDecoder,
+    TransformerDecoderLayer,
+)
+from torchmultimodal.modules.layers.attention import SelfAttention
+from torchmultimodal.modules.layers.position_embedding import (
+    BroadcastedPositionEmbedding,
+)
+from torchmultimodal.utils.common import load_module_from_url
+from torchtext.transforms import CharBPETokenizer
+
+
+PRETRAINED_TOKENIZER_ENCODER_URL = "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/tokenizer-coinrun_1024_encoder.json"
+PRETRAINED_TOKENIZER_MERGES_URL = "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/tokenizer-coinrun_1024_merges.txt"
+PRETRAINED_TEXT_VIDEO_GPT_URL_MAPPING = {
+    "mugen_L32": "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/text_video_gpt_L32_weights-17db9549.pth",
+    "mugen_L16": "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/text_video_gpt_L16_weights-5dfc5a0a.pth",
+    "mugen_L8": "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/text_video_gpt_L8_weights-72b6d2ab.pth",
+}
+
+
+def text_video_gpt(
+    text_seq_len: int = 128,
+    video_seq_len: int = 32,
+    resolution: int = 256,
+    downsample: Tuple[int, int, int] = (4, 32, 32),
+    d_model: int = 768,
+    n_head: int = 8,
+    dropout: float = 0.2,
+    attn_dropout: float = 0.3,
+    num_decoder_layers: int = 12,
+    use_gpt_init: bool = True,
+    pretrained_text_tokenizer_encoder_url: str = PRETRAINED_TOKENIZER_ENCODER_URL,
+    pretrained_text_tokenizer_merges_url: str = PRETRAINED_TOKENIZER_MERGES_URL,
+    pretrained_video_vqvae_model_key: Optional[str] = None,
+    pretrained_text_video_gpt_model_key: Optional[str] = None,
+) -> MultimodalGPT:
+    """Builds a text-to-video GPT model from user inputs
+
+    Parameter defaults follow MUGEN project:
+        * Video VQVAE: https://github.com/mugen-org/MUGEN_baseline/tree/main/generation/experiments/vqvae
+        * GPT: https://github.com/mugen-org/MUGEN_baseline/blob/main/lib/models/gpt/gpt.py#L252
+
+    Args:
+        text_seq_len (int): Length of text sequences after padding. Defaults to ``128``.
+        video_seq_len (int): Length of video sequences sampled from the dataset. Defaults to ``32``. Other
+            values used by MUGEN are ``8``, ``16``.
+        resolution (int): Resolution of the sampled video sequences defining height and width of each frame.
+            Defaults to ``256``.
+        downsample (Tuple[int, int, int]): Ratio by which to disperse along each dimension the sampled sequences.
+            For example, if the original frame is ``(32, 256, 256)``, after downsampling by ``(4, 32, 32)`` the
+            new frame will be of shape ``(8, 8, 8)`` with each dim divided by the rate of downsample. Defaults to
+            ``(4, 32, 32)``.
+        d_model (int): Dimension of the underlying transformer decoder.
+            See :py:class:`torchmultimodal.models.video_gpt.gpt.TransformerDecoderLayer`. Defaults to ``768``.
+        n_head (int): Number of attention heads used by the transformer decoder. Defaults to ``8``.
+        dropout (float): Dropout probability used by the projection layer of the transformer decoder.
+            Defaults to ``0.2``.
+        attn_dropout (float): Dropout probability used by the attention layer of the transformer decoder.
+            Defaults to ``0.3``.
+        num_decoder_layers (int): Number of transformer decoder layers. Defaults to ``12``.
+        use_gpt_init (bool): Whether uses parameter initialization of GPT model. Defaults to ``True``.
+        pretrained_text_tokenizer_encoder_url (str): Remote location of the pretrained text tokenizer encoder file.
+            Defaults to `"MUGEN pretrained tokenizer encoder file
+            "<https://pytorch.s3.amazonaws.com/models/multimodal/mugen/tokenizer-coinrun_1024_encoder.json>`_.
+        pretrained_text_tokenizer_merges_url (str): Remote location of the pretrained text tokenizer merges file.
+            Defaults to `"MUGEN pretrained tokenizer merges file
+            "<https://pytorch.s3.amazonaws.com/models/multimodal/mugen/tokenizer-coinrun_1024_merges.txt>`_.
+        pretrained_video_vqvae_model_key (str, optional): Key to select the pretrained MUGEN VideoVQVAE weights
+            file. For allowed values, see :py:module:`examples/mugen/generation/video_vqvae.py`.
+            Defaults to ``None``.
+        pretrained_text_video_gpt_model_key (str, optional): Key to select the pretrained MUGEN TextVideoGPT
+            weights file. The provided key should match that of MUGEN VideoVQVAE to ensure the two models were
+            pretrained for the same video sequence length. For example ``L32`` means the video sequence length
+            is ``32``. The loaded weights will override those from the frozen VideoVQVAE model.
+            Defaults to ``None``.
+
+    Returns:
+        An instance of :py:class:`torchmultimodal.models.video_gpt.gpt.MultimodalGPT`.
+    """
+
+    # builds text tokenizer from pre-trained
+    tokenizer = CharBPETokenizer(
+        bpe_encoder_path=pretrained_text_tokenizer_encoder_url,
+        bpe_merges_path=pretrained_text_tokenizer_merges_url,
+        unk_token="[UNK]",
+        special_tokens=["[PAD]", "[CLS]", "[SEP]", "[UNK]", "[MASK]"],
+    )
+
+    # builds text tokenizer
+    text_tokenizer = TextTokenizer(
+        context_len=text_seq_len,
+        d_model=d_model,
+        tokenizer=tokenizer,
+    )
+    num_text_tokens = text_tokenizer.num_text_tokens
+
+    # builds video tokenizer
+    video_vqvae = video_vqvae_mugen(
+        pretrained_model_key=pretrained_video_vqvae_model_key,
+        freeze_model=True,
+    )
+    video_vqvae.eval()
+    num_video_tokens = video_vqvae.num_embeddings  # size of the codebook
+
+    # derives the expected latent shape from video input shape
+    video_input_shape = (video_seq_len, resolution, resolution)
+    video_latent_shape = latent_shape(video_input_shape, downsample)
+    video_vqvae_latent_shape = video_vqvae.latent_shape(video_input_shape)
+    # video vqvae will apply convolutions to the input shape which effectively
+    # reduces the size by ``dim//stride`` after each layer
+    # sanity check that the expected and actual latent shapes are consistent
+    if video_latent_shape != video_vqvae_latent_shape:
+        raise ValueError(
+            f"Latent shape derived from video inputs: {video_latent_shape} "
+            f"does not match that of video vqvae: {video_vqvae_latent_shape}"
+        )
+
+    # builds text embedding projection: text_emb is already of output shape `d_model`
+    # generally a projection layer is needed to bridge the tokenizer and
+    # `torchmultimodal.models.gpt.MultimodalTransformerDecoder`, see `video_projection`
+    text_projection = nn.Identity()
+
+    # builds video embedding projection
+    video_projection = nn.Linear(video_vqvae.embedding_dim, d_model, bias=False)
+
+    # builds multimodal decoder
+    text_pos_emb = nn.Embedding(text_seq_len, d_model)
+    video_pos_emb = BroadcastedPositionEmbedding(video_latent_shape, d_model)
+    attention_layer = SelfAttention(attn_dropout=attn_dropout)
+    decoder_layer = TransformerDecoderLayer(
+        d_model, n_head, dropout, attn_module=attention_layer
+    )
+    decoder = TransformerDecoder(decoder_layer, num_decoder_layers)
+    right_shift = RightShift(d_model)
+    mm_decoder = MultimodalTransformerDecoder(
+        text_pos_emb, video_pos_emb, decoder, right_shift
+    )
+
+    model = MultimodalGPT(
+        d_model=d_model,
+        num_in_tokens=num_text_tokens,
+        num_out_tokens=num_video_tokens,
+        latent_shape=video_latent_shape,
+        in_tokenizer=text_tokenizer,
+        out_tokenizer=video_vqvae,
+        mm_decoder=mm_decoder,
+        in_projection=text_projection,
+        out_projection=video_projection,
+        use_gpt_init=use_gpt_init,
+    )
+
+    if pretrained_text_video_gpt_model_key is not None:
+        if (
+            pretrained_text_video_gpt_model_key
+            not in PRETRAINED_TEXT_VIDEO_GPT_URL_MAPPING
+        ):
+            raise KeyError(
+                f"Invalid pretrained model key: {pretrained_text_video_gpt_model_key}"
+            )
+
+        load_module_from_url(
+            model,
+            PRETRAINED_TEXT_VIDEO_GPT_URL_MAPPING[pretrained_text_video_gpt_model_key],
+        )
+
+    return model
+
+
+def latent_shape(
+    input_shape: Tuple[int, ...], downsample: Tuple[int, ...]
+) -> Tuple[int, ...]:
+    """Derives latent shape of video inputs after VQ-VAE encoding"""
+    return tuple([s // d for s, d in zip(input_shape, downsample)])
+
+
+class TextTokenizer(nn.Module):
+    """Converts between text and tokens / embedings
+
+    Wrapper around the tokenizer to be consistent with the API required by
+    :py:class:`torchmultimodal.models.video_gpt.gpt.MultimodalGPT`. It also contains the
+    embedding layer to enable lookup by token ids.
+    """
+
+    def __init__(
+        self,
+        context_len: int,
+        d_model: int,
+        tokenizer: nn.Module,
+    ) -> None:
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.pad_id = self.tokenizer.encode("[PAD]")[0]  # type: ignore
+        self.vocab_size = self.tokenizer.vocab_size  # type: ignore
+        self.context_len = context_len
+        # MUGEN treats padding as unique ids so adding them to the total text tokens
+        # https://github.com/mugen-org/MUGEN_baseline/blob/main/lib/models/gpt/gpt.py#L44
+        self.num_text_tokens = self.vocab_size + context_len
+        self.embedding = nn.Embedding(self.num_text_tokens, d_model)
+
+    def text_to_tokens(self, sentences: List[str]) -> Tensor:
+        """Pads the sentences to be of equal lengths"""
+        tokens = [
+            self.tokenizer.encode(sentence.strip().lower() + " [SEP]")  # type: ignore
+            for sentence in sentences
+        ]
+        token_ids = [t[: self.context_len] for t in tokens]
+        # pad each sentence to be of length `context_len`
+        for i, t in enumerate(token_ids):
+            t += [self.pad_id] * (self.context_len - len(t))
+            token_ids[i] = t
+
+        return torch.Tensor(token_ids).type(torch.int64)
+
+    def encode(self, sentences: List[str], device: str) -> Tensor:
+        """Encodes sentences to token ids"""
+        token_ids = self.text_to_tokens(sentences).to(device)
+        # bump padding token ids by vocab_size so that they do not coincide with un-padded token ids
+        # and that the padding token ids themselves are unique
+        unique_pad_ids = torch.arange(self.context_len, device=device) + self.vocab_size
+        token_ids = torch.where(token_ids == self.pad_id, unique_pad_ids, token_ids)
+        return token_ids
+
+    def _filter_token_ids(self, token_ids: List[int]) -> List[Optional[int]]:
+        """Filters out token ids out side of vocab"""
+        return [
+            token_id
+            for token_id in token_ids
+            if token_id > 0 and token_id <= self.vocab_size
+        ]
+
+    def decode(self, token_ids: Tensor) -> List[str]:
+        """Decodes token ids back to sentences"""
+        sentences = []
+        for _token_ids in token_ids:  # iterate over batches
+            _token_ids = self._filter_token_ids(_token_ids.tolist())
+            sentence = self.tokenizer.decode(_token_ids)  # type: ignore
+            sentences.append(sentence)
+
+        return sentences
+
+    def lookup(self, token_ids: Tensor) -> Tensor:
+        return self.embedding(token_ids)

multimodal/examples/mugen/generation/video_vqvae.py

@@ -0,0 +1,113 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+from torchmultimodal.models.video_gpt.video_vqvae import (
+    preprocess_int_conv_params,
+    VideoDecoder,
+    VideoEncoder,
+)
+
+from torchmultimodal.models.vqvae import VQVAE
+from torchmultimodal.utils.common import load_module_from_url, remove_grad
+
+
+MUGEN_PRETRAINED_MAPPING = {
+    "mugen_L32": "https://download.pytorch.org/models/multimodal/mugen/mugen_video_vqvae_L32.pt",
+    "mugen_L16": "https://download.pytorch.org/models/multimodal/mugen/mugen_video_vqvae_L16.pt",
+    "mugen_L8": "https://download.pytorch.org/models/multimodal/mugen/mugen_video_vqvae_L8.pt",
+}
+
+
+def video_vqvae_mugen(
+    in_channel_dim: int = 3,
+    encoder_hidden_dim: int = 240,
+    encoder_kernel_size: int = 3,
+    n_res_layers: int = 4,
+    attn_hidden_dim: int = 240,
+    num_embeddings: int = 2048,
+    embedding_dim: int = 256,
+    decoder_hidden_dim: int = 240,
+    decoder_kernel_size: int = 3,
+    pretrained_model_key: Optional[str] = None,
+    freeze_model: bool = False,
+) -> VQVAE:
+    """Constructor for MUGEN's Video VQVAE. Expects input video data of shape ``{8,16,32}x256x256``.
+    Trained for tokenization of video data and use in video-audio-text retrieval and generation tasks.
+    See Hayes et al. 2022 for more details: https://arxiv.org/pdf/2204.08058.pdf
+    Code ref:
+    https://github.com/mugen-org/MUGEN_baseline/blob/main/lib/models/video_vqvae/vqvae.py
+    https://github.com/mugen-org/MUGEN_baseline/blob/main/generation/experiments/vqvae/VideoVQVAE_L32.sh
+
+    Args:
+        in_channel_dim (int, optional): Size of channel dim in input. Defaults to ``3``.
+        encoder_hidden_dim (int, optional): Size of channel dims in encoder conv layers. Defaults to ``240``.
+        encoder_kernel_size (int, optional): Kernel size for encoder. Defaults to ``3``.
+        n_res_layers (int, optional): Number of ``AttentionResidualBlocks`` to include in encoder and decoder.
+            Defaults to ``4``.
+        attn_hidden_dim (int, optional): Size of hidden dim of
+            :class:`~torchmultimodal.models.video_gpt.video_vqvae.AttentionResidualBlocks`. Defaults to ``240``.
+        num_embeddings (int, optional): Number of embedding vectors used in
+            :class:`~torchmultimodal.modules.layers.codebook.Codebook`. Defaults to ``2048``.
+        embedding_dim (int, optional): Dimensionality of embedding vectors in
+            :class:`~torchmultimodal.modules.layers.codebook.Codebook`. Defaults to ``256``.
+        decoder_hidden_dim (int, optional): Size of channel dims in decoder conv tranpose layers.
+            Defaults to ``240``.
+        decoder_kernel_size (int, optional): Kernel size for decoder. Defaults to ``3``.
+        pretrained_model_key (str, optional): Load a specified MUGEN VQVAE checkpoint.
+        freeze_model (bool): Whether to freeze the weights of the pretrained model. Defaults to ``False``.
+
+    Returns:
+        An instance of :class:`~torchmultimodal.models.vqvae.VQVAE` constructed with:
+            * :class:`~torchmultimodal.model.video_gpt.video_vqvae.VideoEncoder`
+            * :class:`~torchmultimodal.model.video_gpt.video_vqvae.VideoDecoder`
+    """
+    encoder_strides = ((2, 2, 2), (2, 2, 2), (1, 2, 2), (1, 2, 2), (1, 2, 2), (1, 1, 1))
+    decoder_strides = ((2, 2, 2), (2, 2, 2), (1, 2, 2), (1, 2, 2), (1, 2, 2))
+    encoder_n_layers = len(encoder_strides)
+    decoder_n_layers = len(decoder_strides)
+    encoder_in_channel_dims = (in_channel_dim,) + (encoder_hidden_dim,) * max(
+        encoder_n_layers - 1, 0
+    )
+    decoder_out_channel_dims = (decoder_hidden_dim,) * max(decoder_n_layers - 1, 0) + (
+        in_channel_dim,
+    )
+    encoder_kernel_sizes_fixed = preprocess_int_conv_params(
+        encoder_in_channel_dims, encoder_kernel_size
+    )
+    decoder_kernel_sizes_fixed = preprocess_int_conv_params(
+        decoder_out_channel_dims, decoder_kernel_size
+    )
+
+    encoder = VideoEncoder(
+        encoder_in_channel_dims,
+        encoder_kernel_sizes_fixed,
+        encoder_strides,
+        embedding_dim,
+        n_res_layers,
+        attn_hidden_dim,
+    )
+    decoder = VideoDecoder(
+        decoder_out_channel_dims,
+        decoder_kernel_sizes_fixed,
+        decoder_strides,
+        embedding_dim,
+        n_res_layers,
+        attn_hidden_dim,
+    )
+    model = VQVAE(encoder, decoder, num_embeddings, embedding_dim)
+
+    if pretrained_model_key is not None:
+        if pretrained_model_key not in MUGEN_PRETRAINED_MAPPING.keys():
+            raise KeyError(f"Invalid pretrained model key: {pretrained_model_key}")
+
+        load_module_from_url(model, MUGEN_PRETRAINED_MAPPING[pretrained_model_key])
+
+        if freeze_model:
+            remove_grad(model)
+
+    return model

multimodal/examples/mugen/retrieval/README.md

@@ -0,0 +1,34 @@
+# MUGEN Retrieval
+
+This directory contains reference training and evaluation scripts for MUGEN's video-text retrieval model, including a tutorial notebook for the model usage [Colab](https://colab.research.google.com/drive/1gZfz1jsy79CNCK9t2_r43yt3z7v-w4HS?usp=sharing) or [GitHub](https://github.com/facebookresearch/multimodal/blob/main/examples/mugen/retrieval/tutorial.ipynb).
+
+## Model
+MUGEN's video-text retrieval model follows from [VideoCLIP](https://arxiv.org/abs/2109.14084), a contrastive model for video and text.
+
+The name "VideoCLIP" refers to its similarities to OpenAI's [CLIP](https://arxiv.org/abs/2103.00020), which was originally proposed for zero-shot learning of image classification tasks by “drawing cues” from text data with the corresponding visual concepts. Unlike various predecessor models based on supervised learning, CLIP does not have to be trained on the task-specific datasets or fine-tuned with a task-specific head. The model learns a joint embedding space for both image and text data and optimizes a scaled cosine similarity function between the image and text embedding vectors. The loss function is the sum of the normalized cosine similarities for every pair of image-and-text samples. Each embedding is trained with a unimodal encoder, e.g., a transformer for text, vision transformer (ViT) or ResNet for image.
+
+The VideoCLIP model follows the CLIP architecture but replaces the image encoder with a video encoder. VideoCLIP's video encoder is backed by [Separable 3D CNN (S3D)](https://arxiv.org/abs/1712.04851), a video classification model, and the text encoder is backed by [DistilBERT](https://arxiv.org/abs/1910.01108), a lightweight transformer for language modeling.
+
+## Training
+The configurable parameters for training can be found in `configs/train.yaml`. Note that the training script supports training on 1 or more devices on a single node. Then run the following command:
+```
+python train.py config=configs/train.yaml
+```
+A checkpoint file with the best-performing weights will be saved under `{default_root_dir}/lightning_logs/`, where `default_root_dir` is specified in the training config. If `default_root_dir` is `null`, then it will act as your working directory.
+
+## Evaluation
+The configurable parameters for evaluation can be found in `configs/eval.yaml`. You can choose to replace `checkpoint_path` with the path to your checkpoint from the training step, or keep the default `checkpoint_path` to load the MUGEN authors' weights (fit to our implementation). Then run the following command:
+```
+python eval.py config=configs/eval.yaml
+```
+
+Using the default arguments in `configs/eval.yaml` (including the MUGEN authors' published weights), we ran the evaluation script on the full MUGEN test set and got the following results:
+
+| Metric (%)                | MUGEN Results | TorchMultimodal Results   |
+| -----------               | -----------   | -----------               |
+| Text2video top-1 recall   | 8.54          | 8.26                      |
+| Text2video top-5 recall   | 22.50         | 22.34                     |
+| Text2video top-10 recall  | 31.71         | 31.68                     |
+| Video2text top-1 recall   | 10.61         | 10.79                     |
+| Video2text top-5 recall   | 25.72         | 25.70                     |
+| Video2text top-10 recall  | 34.70         | 34.60                     |

multimodal/examples/mugen/retrieval/configs/eval.yaml

@@ -0,0 +1,48 @@
+_target_: examples.mugen.retrieval.definitions.EvaluationArgs
+dataset_args:
+  _target_: examples.mugen.data.mugen_dataset.MUGENDatasetArgs
+  data_path: "datasets/coinrun/coinrun_dataset_jsons/release"
+  asset_path: "datasets/coinrun/assets"
+  sample_every_n_frames: 3
+  sequence_length: 32
+  audio_sample_rate: 22050
+  audio_sample_length: 70560
+  resolution: 256
+  bbox_smap_for_agent: False
+  bbox_smap_for_monsters: False
+  use_manual_annotation: True
+  use_auto_annotation: False
+  use_downsampled_trainset: False
+  fixed_start_idx: False
+  get_game_frame: True
+  get_seg_map: False
+  get_text_desc: True
+  get_audio: False
+  debug: False
+datamodule_args:
+  _target_: examples.mugen.retrieval.definitions.DataModuleArgs
+  batch_size: 16
+  num_workers: 4
+  shuffle: False
+  bert_text_transform:
+    _target_: examples.mugen.retrieval.definitions.BertTextTransformArgs
+  video_transform:
+    _target_: examples.mugen.retrieval.definitions.VideoTransformArgs
+lightningmodule_args:
+  _target_: examples.mugen.retrieval.definitions.LightningModuleArgs
+  logit_scale: 0.07
+  logit_scale_max: 100.0
+videoclip_args:
+  _target_: examples.mugen.retrieval.definitions.VideoCLIPArgs
+  text_pretrained: False
+  text_trainable: False
+  text_model_name: "distilbert-base-uncased"
+  text_model_config: null
+  text_padding_value: 0
+  video_pretrained: False
+  video_trainable: False
+  video_pretrain_path: "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/S3D_kinetics400.pt"
+  proj_out_dim: 256
+  proj_dropout: 0.1
+checkpoint_path: "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/videoclip_lightning_mugen.pt"
+accelerator: "auto"

multimodal/examples/mugen/retrieval/configs/train.yaml

@@ -0,0 +1,53 @@
+_target_: examples.mugen.retrieval.definitions.TrainingArgs
+dataset_args:
+  _target_: examples.mugen.data.mugen_dataset.MUGENDatasetArgs
+  data_path: "datasets/coinrun/coinrun_dataset_jsons/release"
+  asset_path: "datasets/coinrun/assets"
+  sample_every_n_frames: 3
+  sequence_length: 32
+  audio_sample_rate: 22050
+  audio_sample_length: 70560
+  resolution: 224
+  bbox_smap_for_agent: False
+  bbox_smap_for_monsters: False
+  use_manual_annotation: True
+  use_auto_annotation: False
+  use_downsampled_trainset: False
+  fixed_start_idx: False
+  get_game_frame: True
+  get_seg_map: False
+  get_text_desc: True
+  get_audio: False
+  debug: False
+datamodule_args:
+  _target_: examples.mugen.retrieval.definitions.DataModuleArgs
+  batch_size: 16
+  num_workers: 4
+  shuffle: False
+  bert_text_transform:
+    _target_: examples.mugen.retrieval.definitions.BertTextTransformArgs
+  video_transform:
+    _target_: examples.mugen.retrieval.definitions.VideoTransformArgs
+lightningmodule_args:
+  _target_: examples.mugen.retrieval.definitions.LightningModuleArgs
+  logit_scale: 0.07
+  logit_scale_max: 100.0
+  learning_rate: 0.001
+  weight_decay: 0.001
+videoclip_args:
+  _target_: examples.mugen.retrieval.definitions.VideoCLIPArgs
+  text_pretrained: True
+  text_trainable: False
+  text_model_name: "distilbert-base-uncased"
+  text_model_config: null
+  text_padding_value: 0
+  video_pretrained: True
+  video_trainable: True
+  video_pretrain_path: "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/S3D_kinetics400.pt"
+  proj_out_dim: 256
+  proj_dropout: 0.1
+accelerator: "auto"
+devices: 4
+max_epochs: 20
+log_every_n_steps: 100
+default_root_dir: null

multimodal/examples/mugen/retrieval/definitions.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+from examples.mugen.data.mugen_dataset import MUGENDatasetArgs
+
+from torchmultimodal.transforms.video_transform import (
+    DEFAULT_MEAN,
+    DEFAULT_RESIZE_SHAPE,
+    DEFAULT_STD,
+    MUGEN_DEFAULT_TIME_SAMPLES,
+)
+
+
+@dataclass
+class BertTextTransformArgs:
+    vocab_file: str = "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt"
+    do_lower_case: bool = True
+    start_token: int = 101
+    end_token: int = 102
+    padding_value: int = 0
+
+
+@dataclass
+class VideoTransformArgs:
+    time_samples: int = MUGEN_DEFAULT_TIME_SAMPLES
+    mean: Tuple[float] = DEFAULT_MEAN
+    std: Tuple[float] = DEFAULT_STD
+    resize_shape: Tuple[int, int] = DEFAULT_RESIZE_SHAPE
+
+
+@dataclass
+class DataModuleArgs:
+    batch_size: int = 16
+    num_workers: int = 4
+    shuffle: bool = False
+    bert_text_transform: BertTextTransformArgs = BertTextTransformArgs()
+    video_transform: VideoTransformArgs = VideoTransformArgs()
+
+
+@dataclass
+class LightningModuleArgs:
+    logit_scale: float = 0.07
+    logit_scale_max: float = 100.0
+    learning_rate: float = 1e-3
+    weight_decay: float = 1e-3
+    recall_ks: Tuple[int] = (1, 5, 10)
+
+
+@dataclass
+class VideoCLIPArgs:
+    text_pretrained: bool = False
+    text_trainable: bool = False
+    text_model_name: str = "distilbert-base-uncased"
+    text_model_config: Optional[Dict[str, Any]] = None
+    text_padding_value: int = 0
+    video_pretrained: bool = False
+    video_trainable: bool = False
+    video_pretrain_path: str = (
+        "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/S3D_kinetics400.pt"
+    )
+    proj_out_dim: int = 256
+    proj_dropout: float = 0.1
+
+
+@dataclass
+class EvaluationArgs:
+    dataset_args: MUGENDatasetArgs = MUGENDatasetArgs(
+        get_game_frame=True,
+        get_text_desc=True,
+        resolution=256,
+        fixed_start_idx=False,
+        use_manual_annotation=True,
+        use_auto_annotation=False,
+    )
+    datamodule_args: DataModuleArgs = DataModuleArgs()
+    lightningmodule_args: LightningModuleArgs = LightningModuleArgs()
+    videoclip_args: VideoCLIPArgs = VideoCLIPArgs()
+    checkpoint_path: str = "https://pytorch.s3.amazonaws.com/models/multimodal/mugen/videoclip_lightning_mugen.pt"
+    accelerator: str = "auto"
+
+
+@dataclass
+class TrainingArgs:
+    dataset_args: MUGENDatasetArgs = MUGENDatasetArgs(
+        get_game_frame=True,
+        get_text_desc=True,
+        resolution=224,
+        fixed_start_idx=False,
+        use_manual_annotation=True,
+        use_auto_annotation=False,
+    )
+    datamodule_args: DataModuleArgs = DataModuleArgs()
+    lightningmodule_args: LightningModuleArgs = LightningModuleArgs()
+    videoclip_args: VideoCLIPArgs = VideoCLIPArgs()
+    accelerator: str = "auto"
+    devices: int = 4
+    max_epochs: int = 1000
+    log_every_n_steps: int = 100
+    default_root_dir: Optional[str] = None

multimodal/examples/mugen/retrieval/eval.py

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from examples.mugen.data.bert_text_transform import BertTextTransform
+from examples.mugen.data.mugen_datamodules import MUGENDataModule
+from examples.mugen.data.mugen_dataset import MUGENDatasetArgs
+from examples.mugen.retrieval.model import VideoCLIPLightningModule
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+from torchmultimodal.transforms.video_transform import VideoTransform
+
+
+def get_yaml_config():
+    cli_conf = OmegaConf.from_cli()
+    if "config" not in cli_conf:
+        raise ValueError(
+            "Please pass 'config' to specify configuration yaml file for running VideoCLIP evaluation"
+        )
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = instantiate(yaml_conf)
+    return conf
+
+
+def evaluate():
+    args = get_yaml_config()
+
+    dataset_args: MUGENDatasetArgs = args.dataset_args
+    datamodule = MUGENDataModule(
+        dataset_args,
+        text_transform=BertTextTransform(
+            **vars(args.datamodule_args.bert_text_transform)
+        ),
+        video_transform=VideoTransform(**vars(args.datamodule_args.video_transform)),
+        batch_size=args.datamodule_args.batch_size,
+        num_workers=args.datamodule_args.num_workers,
+        shuffle=args.datamodule_args.shuffle,
+    )
+
+    model = VideoCLIPLightningModule.load_from_checkpoint(
+        args.checkpoint_path,
+        **vars(args.lightningmodule_args),
+        **vars(args.videoclip_args),
+    )
+
+    trainer = Trainer(accelerator=args.accelerator, devices=1)
+    trainer.test(model, dataloaders=datamodule.test_dataloader())
+
+
+if __name__ == "__main__":
+    evaluate()

multimodal/examples/mugen/retrieval/model.py

@@ -0,0 +1,145 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import warnings
+from typing import Any, Tuple
+
+import torch
+
+from examples.mugen.retrieval.video_clip import videoclip
+from pytorch_lightning import LightningModule
+from torchmetrics import Recall
+
+from torchmultimodal.modules.losses.contrastive_loss_with_temperature import (
+    ContrastiveLossWithTemperature,
+)
+
+
+class VideoCLIPLightningModule(LightningModule):
+    """PyTorch Lightning module for evaluating VideoCLIP model.
+    Args:
+        logit_scale (float): Initial log-temperature value for contrastive loss funtion.
+            Defaults to ``0.07``, MUGEN's log-temperature value at initialization.
+        logit_scale_max (float): Maximum log-temperature value for contrastive loss function.
+            Defaults to ``100``, MUGEN's maximum log-temperature value.
+        learning_rate (float): optimizer learning rate.
+            Defaults to ``1e-3``, MUGEN's learning rate.
+        weight_decay (float): optimizer weight decay.
+            Defaults to ``1e-3``, MUGEN's weight decay.
+        recall_ks (Tuple[int]): tuple of top-``k``'s for calculating recall.
+            Defaults to ``(1, 5, 10)``, i.e. top-1 recall, top-5 recall, and top-10 recall.
+        **videoclip_kwargs (Any): Keyword arguments for the videoCLIP model builder.
+    """
+
+    def __init__(
+        self,
+        logit_scale: float = 0.07,
+        logit_scale_max: float = 100,
+        learning_rate: float = 1e-3,
+        weight_decay: float = 1e-3,
+        recall_ks: Tuple[int] = (1, 5, 10),
+        **videoclip_kwargs: Any,
+    ):
+        super().__init__()
+        self.model = videoclip(**videoclip_kwargs)
+        self.contrastive_loss = ContrastiveLossWithTemperature(
+            logit_scale=logit_scale,
+            logit_scale_min=None,
+            logit_scale_max=logit_scale_max,
+        )
+        self.lr = learning_rate
+        self.weight_decay = weight_decay
+
+        self.recall_ks = set(recall_ks)
+        if len(self.recall_ks) != len(recall_ks):
+            warnings.warn("Duplicate `k` values in `recall_ks` are ignored.")
+        self.metrics = torch.nn.ModuleDict()
+        for k in self.recall_ks:
+            self.metrics.update(
+                {f"v2t_recall_{k}": Recall(top_k=k), f"t2v_recall_{k}": Recall(top_k=k)}
+            )
+
+    def _collect_embeddings(self, outputs):
+        text_embeddings = [batch.embeddings_a for batch in outputs]
+        video_embeddings = [batch.embeddings_b for batch in outputs]
+
+        embeddings = {
+            "text": torch.cat(text_embeddings),
+            "video": torch.cat(video_embeddings),
+        }
+        return embeddings
+
+    def _compute_recall(self, split, text_embedding, video_embedding):
+        similarity_matrix = text_embedding @ video_embedding.T
+        num_samples = similarity_matrix.shape[0]
+        target_matrix = torch.eye(
+            n=num_samples, dtype=int, device=similarity_matrix.device
+        )
+
+        for k in self.recall_ks:
+            v2t_recall = self.metrics[f"v2t_recall_{k}"]
+            v2t_recall(preds=similarity_matrix.T, target=target_matrix)
+            self.log(f"{split}/Recall@{k} (video query, text retrieval)", v2t_recall)
+
+            t2v_recall = self.metrics[f"t2v_recall_{k}"]
+            t2v_recall(preds=similarity_matrix, target=target_matrix)
+            self.log(f"{split}/Recall@{k} (text query, video retrieval)", t2v_recall)
+
+    def configure_optimizers(self):
+        params = self.parameters()
+        optimizer = torch.optim.AdamW(
+            params, lr=self.lr, weight_decay=self.weight_decay
+        )
+        return optimizer
+
+    def training_step(self, batch, batch_idx):
+        text, video = batch.get("text"), batch.get("video")
+        model_output = self.model(features_a=text, features_b=video)
+        loss = self.contrastive_loss(
+            model_output.embeddings_a, model_output.embeddings_b
+        )
+        self.log(
+            "train/loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
+        )
+        return {"loss": loss, "model_output": model_output}
+
+    def validation_step(self, batch, batch_idx):
+        text, video = batch.get("text"), batch.get("video")
+        model_output = self.model(features_a=text, features_b=video)
+        loss = self.contrastive_loss(
+            model_output.embeddings_a, model_output.embeddings_b
+        )
+        self.log(
+            "validation/loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            logger=True,
+        )
+        return {"loss": loss, "model_output": model_output}
+
+    def validation_epoch_end(self, outputs):
+        model_outputs = [batch["model_output"] for batch in outputs]
+        all_embeddings = self._collect_embeddings(model_outputs)
+        text_embedding, video_embedding = (
+            all_embeddings["text"],
+            all_embeddings["video"],
+        )
+        self._compute_recall("validation", text_embedding, video_embedding)
+
+    def test_step(self, batch, batch_idx):
+        text, video = batch.get("text"), batch.get("video")
+        model_output = self.model(features_a=text, features_b=video)
+        return model_output
+
+    def test_epoch_end(self, outputs):
+        all_embeddings = self._collect_embeddings(outputs)
+        text_embedding, video_embedding = (
+            all_embeddings["text"],
+            all_embeddings["video"],
+        )
+        self._compute_recall("test", text_embedding, video_embedding)

multimodal/examples/mugen/retrieval/train.py

@@ -0,0 +1,67 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from examples.mugen.data.bert_text_transform import BertTextTransform
+from examples.mugen.data.mugen_datamodules import MUGENDataModule
+from examples.mugen.data.mugen_dataset import MUGENDatasetArgs
+from examples.mugen.retrieval.model import VideoCLIPLightningModule
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
+from torchmultimodal.transforms.video_transform import VideoTransform
+
+
+def get_yaml_config():
+    cli_conf = OmegaConf.from_cli()
+    if "config" not in cli_conf:
+        raise ValueError(
+            "Please pass 'config' to specify configuration yaml file for running VideoCLIP training"
+        )
+    yaml_conf = OmegaConf.load(cli_conf.config)
+    conf = instantiate(yaml_conf)
+    return conf
+
+
+def train():
+    args = get_yaml_config()
+
+    dataset_args: MUGENDatasetArgs = args.dataset_args
+    datamodule = MUGENDataModule(
+        dataset_args,
+        text_transform=BertTextTransform(
+            **vars(args.datamodule_args.bert_text_transform)
+        ),
+        video_transform=VideoTransform(**vars(args.datamodule_args.video_transform)),
+        batch_size=args.datamodule_args.batch_size,
+        num_workers=args.datamodule_args.num_workers,
+        shuffle=args.datamodule_args.shuffle,
+    )
+
+    model = VideoCLIPLightningModule(
+        **vars(args.lightningmodule_args),
+        **vars(args.videoclip_args),
+    )
+
+    checkpoint_callback = ModelCheckpoint(save_top_k=-1)
+    trainer = Trainer(
+        accelerator=args.accelerator,
+        devices=args.devices,
+        strategy="ddp_find_unused_parameters_false",
+        max_epochs=args.max_epochs,
+        log_every_n_steps=args.log_every_n_steps,
+        default_root_dir=args.default_root_dir,
+        callbacks=[checkpoint_callback],
+    )
+    trainer.fit(
+        model=model,
+        train_dataloaders=datamodule.train_dataloader(),
+        val_dataloaders=datamodule.val_dataloader(),
+    )
+
+
+if __name__ == "__main__":
+    train()