[feat] Add UNITER model wrapper

Add UNITER model to mmf registry with support
for pretraining through yaml head configs.

ghstack-source-id: 028fd3e
Pull Request resolved: #1127

Author: Ryan-Qiyu-Jiang

mmf/configs/models/uniter/defaults.yaml

@@ -0,0 +1,28 @@
+model_config:
+  uniter:
+    heads:
+      vqa2:
+        type: mlp
+        freeze: false
+        lr_multiplier: 1.0
+        in_dim: 768
+        hidden_size: 1536
+        num_labels: 3129
+        pooler_name: bert_pooler
+    text_embeddings:
+      type: bert_embeddings
+    image_embeddings:
+      type: uniter_image_embeddings
+      params:
+        name: 'uniter_image_embeddings'
+    encoder:
+      type: transformer
+      params:
+        bert_model_name: bert-base-uncased
+        hidden_size: 768
+        num_hidden_layers: 12
+        num_attention_heads: 12
+        output_attentions: false
+        output_hidden_states: false
+    tasks:
+    - vqa2

mmf/models/__init__.py

@@ -17,7 +17,7 @@
 from .visual_bert import VisualBERT
 from .vilbert import ViLBERT
 from .albef.vit import AlbefVitEncoder
-
+from .uniter import UNITER
 
 __all__ = [
     "TopDownBottomUp",
@@ -43,4 +43,5 @@
     "UnimodalModal",
     "UnimodalText",
     "AlbefVitEncoder",
+    "UNITER",
 ]

mmf/models/uniter.py

@@ -7,14 +7,16 @@
 import copy
 import logging
 import random
+from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import torch
 from mmf.common.registry import registry
+from mmf.models import BaseModel
 from mmf.modules.losses import MMFLoss
 from mmf.utils.general import retry_n
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import MISSING, DictConfig, OmegaConf
 from torch import Tensor, nn
 from transformers.modeling_bert import BertConfig, BertEmbeddings, BertModel
 
@@ -624,3 +626,150 @@ def _remove_mismatched_captions(self, processed_sample_list: Dict[str, Tensor]):
                 x = x[pos_pairs_mask]
             else:
                 x = x[pos_pairs_mask, ::]
+
+
+@registry.register_model("uniter")
+class UNITER(BaseModel):
+    """ Modification for Joint Vision-Language Encoding
+    """
+
+    @dataclass
+    class Config:
+        random_init: bool = False
+        bert_model_name: str = "bert-base-uncased"
+        img_dim: int = 2048
+        hidden_size: int = 768
+        hidden_dropout_prob: float = 0
+        text_embeddings: Any = field(default_factory=lambda: {})
+        encoder: Any = field(default_factory=lambda: {})
+        heads: Any = MISSING
+        losses: Any = field(default_factory=lambda: {})
+        tasks: Any = MISSING
+        do_pretraining: bool = False
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = OmegaConf.create({**asdict(self.Config()), **config})
+        self.do_pretraining = self.config.do_pretraining
+
+    @classmethod
+    def config_path(cls):
+        return "configs/models/uniter/defaults.yaml"
+
+    def build(self):
+        params = dict(
+            **self.config,
+            head_configs=self.config.heads,
+            loss_configs=self.config.losses,
+        )
+        if self.do_pretraining:
+            self.uniter = UNITERForPretraining(**params)
+        else:
+            self.uniter = UNITERForClassification(**params)
+
+        self.tasks = self.config.tasks
+        if isinstance(self.tasks, str):
+            self.tasks = self.tasks.split(",")
+
+    def init_losses(self):
+        """
+        Defer loss management to submodels,
+        do nothing when called by build_model.
+        """
+
+    def add_pos_feat(self, sample_list: Dict[str, Tensor]):
+        assert "image_info_0" in sample_list
+        assert "bbox" in sample_list["image_info_0"]
+
+        # (x1, y1, x2, y2), dim = (bs, num_feats, 4)
+        bboxs = torch.tensor(sample_list["image_info_0"]["bbox"])[:, :, :4]
+        norm_xy = torch.clone(bboxs)
+        # if bboxs are not normalized, just do it here
+        if norm_xy[0, 0, 0] < 1:
+            img_h = (
+                torch.tensor(sample_list["image_info_0"]["image_height"])
+                .unsqueeze(1)
+                .unsqueeze(1)
+            )  # (bs,)
+            img_w = (
+                torch.tensor(sample_list["image_info_0"]["image_width"])
+                .unsqueeze(1)
+                .unsqueeze(1)
+            )  # (bs,)
+            max_image_size = torch.cat([img_w, img_h, img_w, img_h], dim=-1)
+            max_image_size = max_image_size.to(norm_xy.device)
+            norm_xy /= max_image_size
+
+        bbox_w = (norm_xy[:, :, 2] - norm_xy[:, :, 0]).unsqueeze(-1)
+        bbox_h = (norm_xy[:, :, 3] - norm_xy[:, :, 1]).unsqueeze(-1)
+        area = bbox_w * bbox_h
+        # normalized (x1, y1, x2, y2, w, h, area)
+        pos_feat = torch.cat([norm_xy, bbox_w, bbox_h, area], dim=-1).to(
+            sample_list["image_feature_0"]
+        )
+        sample_list["img_pos_feat"] = pos_feat
+
+    def add_custom_params(self, sample_list: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        image_feat = sample_list["image_feat"] = sample_list["image_feature_0"]
+
+        image_info = getattr(sample_list, "image_info_0", {})
+        image_dim = getattr(image_info, "max_features", None)
+        sample_list["image_dim"] = image_dim
+
+        image_mask = torch.arange(image_feat.size(-2), device=image_feat.device).expand(
+            image_feat.size()[:-1]
+        )
+        if len(image_dim.size()) < len(image_mask.size()):
+            image_dim = image_dim.unsqueeze(-1)
+            assert len(image_dim.size()) == len(image_mask.size())
+        image_mask = image_mask < image_dim
+        sample_list["image_mask"] = image_mask.long()
+
+        attention_mask = torch.cat(
+            (sample_list["input_mask"], sample_list["image_mask"]), dim=-1
+        )
+        sample_list["attention_mask"] = attention_mask
+        task_index = torch.randint(len(self.tasks), (1,)).item()
+        sample_list["task"] = self.tasks[task_index]
+        sample_list["position_ids"] = torch.arange(
+            0,
+            sample_list["input_ids"].size(1),
+            dtype=torch.long,
+            device=image_feat.device,
+        ).unsqueeze(0)
+
+        self.add_pos_feat(sample_list)
+        return sample_list
+
+    def forward(self, sample_list: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        sample_list = self.add_custom_params(sample_list)
+        return self.uniter(sample_list)
+
+    def get_attention_mask(
+        self,
+        sample_list: Dict[str, Tensor],
+        text_embedding: Tensor,
+        image_embedding: Tensor,
+    ) -> Tensor:
+        image_mask = getattr(sample_list, "image_mask", None)
+
+        if image_mask is not None and sample_list.input_mask is not None:
+            attention_mask = torch.cat((sample_list.input_mask, image_mask), dim=-1)
+        elif image_mask is not None:
+            text_mask = torch.ones(
+                text_embedding.size()[:-1],
+                dtype=text_embedding.dtype,
+                device=text_embedding.device,
+            )
+            attention_mask = torch.cat((image_mask, text_mask), dim=-1)
+        elif sample_list.input_mask is not None:
+            image_mask = torch.ones(
+                image_embedding.size()[:-1],
+                dtype=image_embedding.dtype,
+                device=image_embedding.device,
+            )
+            attention_mask = torch.cat((image_mask, sample_list.input_mask), dim=-1)
+        else:
+            attention_mask = None
+
+        return attention_mask

tests/models/test_uniter.py

@@ -1,6 +1,8 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
+import gc
 import unittest
 
+import tests.test_utils as test_utils
 import torch
 from mmf.common.sample import SampleList
 from mmf.models.uniter import (
@@ -9,6 +11,9 @@
     UNITERImageEmbeddings,
     UNITERModelBase,
 )
+from mmf.utils.build import build_model
+from mmf.utils.configuration import Configuration
+from mmf.utils.env import setup_imports, teardown_imports
 from mmf.utils.general import get_current_device
 from omegaconf import OmegaConf
 
@@ -166,3 +171,107 @@ def test_uniter_for_pretraining(self):
 
             self.assertTrue("losses" in model_output)
             self.assertTrue(loss_name in model_output["losses"])
+
+
+class TestUniterModel(unittest.TestCase):
+    def setUp(self):
+        test_utils.setup_proxy()
+        setup_imports()
+        model_name = "uniter"
+        args = test_utils.dummy_args(model=model_name, dataset="vqa2")
+        configuration = Configuration(args)
+        config = configuration.get_config()
+        model_config = config.model_config[model_name]
+        model_config.model = model_name
+        model_config.losses = {"vqa2": "logit_bce"}
+        model_config.do_pretraining = False
+        model_config.tasks = "vqa2"
+        classification_config_dict = {
+            "do_pretraining": False,
+            "tasks": "vqa2",
+            "heads": {"vqa2": {"type": "mlp", "num_labels": 3129}},
+            "losses": {"vqa2": "logit_bce"},
+        }
+        classification_config = OmegaConf.create(
+            {**model_config, **classification_config_dict}
+        )
+
+        pretraining_config_dict = {
+            "do_pretraining": True,
+            "tasks": "wra",
+            "heads": {"wra": {"type": "wra"}},
+        }
+        pretraining_config = OmegaConf.create(
+            {**model_config, **pretraining_config_dict}
+        )
+
+        self.model_for_classification = build_model(classification_config)
+        self.model_for_pretraining = build_model(pretraining_config)
+
+    def tearDown(self):
+        teardown_imports()
+        del self.model_for_classification
+        del self.model_for_pretraining
+        gc.collect()
+
+    def _get_sample_list(self):
+        bs = 8
+        num_feats = 100
+        max_sentence_len = 25
+        img_dim = 2048
+        vqa_cls_dim = 3129
+        input_ids = torch.ones((bs, max_sentence_len), dtype=torch.long)
+        input_mask = torch.ones((bs, max_sentence_len), dtype=torch.long)
+        img_feat = torch.rand((bs, num_feats, img_dim))
+
+        max_features = torch.ones((bs, num_feats)) * num_feats
+        bbox = torch.randint(50, 200, (bs, num_feats, 4)).float()
+        image_height = torch.randint(100, 300, (bs,))
+        image_width = torch.randint(100, 300, (bs,))
+        image_info = {
+            "max_features": max_features,
+            "bbox": bbox,
+            "image_height": image_height,
+            "image_width": image_width,
+        }
+        targets = torch.rand((bs, vqa_cls_dim))
+        is_correct = torch.ones((bs,), dtype=torch.long)
+
+        sample_list = SampleList()
+        sample_list.add_field("input_ids", input_ids)
+        sample_list.add_field("image_feature_0", img_feat)
+        sample_list.add_field("input_mask", input_mask)
+        sample_list.add_field("image_info_0", image_info)
+        sample_list.add_field("targets", targets)
+        sample_list.add_field("is_correct", is_correct)
+        sample_list = sample_list.to(get_current_device())
+        return sample_list
+
+    def test_uniter_for_classification(self):
+        self.model_for_classification.eval()
+        self.model_for_classification = self.model_for_classification.to(
+            get_current_device()
+        )
+        sample_list = self._get_sample_list()
+
+        sample_list.dataset_name = "vqa2"
+        sample_list.dataset_type = "test"
+        with torch.no_grad():
+            model_output = self.model_for_classification(sample_list)
+
+        self.assertTrue("losses" in model_output)
+        self.assertTrue("test/vqa2/logit_bce" in model_output["losses"])
+
+    def test_uniter_for_pretraining(self):
+        self.model_for_pretraining.eval()
+        self.model_for_pretraining = self.model_for_pretraining.to(get_current_device())
+        sample_list = self._get_sample_list()
+        sample_list["tasks"] = "wra"
+
+        sample_list.dataset_name = "vqa2"
+        sample_list.dataset_type = "test"
+        with torch.no_grad():
+            model_output = self.model_for_pretraining(sample_list)
+
+        self.assertTrue("losses" in model_output)
+        self.assertTrue("wra_loss" in model_output["losses"])