[feat] Add UNITER model wrapper (#1127)

Summary:
Pull Request resolved: #1127

Add UNITER BaseModel for training and testing from MMF.
This model defers to either the classification or pretraining
model depending on its config.
For an example config consult the project dir or unit tests.

Test Plan:
### Unit Tests
Tested forward passes for classification and pertaining from mmf basemodel build from config.
### End to End
Model tested end-to-end using butd extracted features on winoground.
Will work on converting caffe feature extraction to pytorch so we can add Uniter and Villa checkpoints usefully.

Reviewed By: ebsmothers

Differential Revision: D31768457

Pulled By: Ryan-Qiyu-Jiang

fbshipit-source-id: b311419f4b1431a2cf8bb5322bd08d80e8a883c3

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

@@ -15,11 +15,11 @@
 from .pythia import Pythia
 from .top_down_bottom_up import TopDownBottomUp
 from .unimodal import UnimodalBase, UnimodalText, UnimodalModal
+from .uniter import UNITER
 from .vilbert import ViLBERT
 from .vilt import ViLT
 from .visual_bert import VisualBERT
 
-
 __all__ = [
     "TopDownBottomUp",
     "Pythia",
@@ -45,4 +45,5 @@
     "UnimodalText",
     "AlbefVitEncoder",
     "ViLT",
+    "UNITER",
 ]

mmf/models/uniter.py

@@ -7,14 +7,16 @@
 import logging
 import random
 from collections import MutableMapping, namedtuple
+from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Tuple, 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
 
@@ -310,8 +312,8 @@ def __init__(
         for task in self.tasks:
             assert task in head_configs, (
                 f"Task {task} is specified in your model configs"
-                + " but there is no head configured for the task."
-                + "Head configs can be added under model_config.heads"
+                + " but there is no head configured for the task. "
+                + "Head configs can be added under model_config.heads "
                 + "in your yaml configs. Either remove this task if UNITER"
                 + " is not meant to run on a dataset named {task}"
                 + " or add a head config."
@@ -603,3 +605,164 @@ 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):
+        configs = dict(**self.config)
+        configs["head_configs"] = configs.pop("heads")
+        configs["loss_configs"] = configs.pop("losses")
+        params_keys = [
+            "head_configs",
+            "loss_configs",
+            "tasks",
+            "random_init",
+            "bert_model_name",
+            "img_dim",
+            "hidden_size",
+            "hidden_dropout_prob",
+            "text_embeddings",
+            "encoder",
+        ]
+        if self.do_pretraining:
+            # take value from config when the key exists,
+            # otherwise use constructor defaults
+            params_keys += ["mask_probability"]
+            params = {key: configs[key] for key in params_keys if key in configs}
+            self.uniter = UNITERForPretraining(**params)
+        else:
+            params = {key: configs[key] for key in params_keys if key in configs}
+            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.
+        """
+        pass
+
+    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()
+
+        sample_list["attention_mask"] = torch.cat(
+            (sample_list["input_mask"], sample_list["image_mask"]), dim=-1
+        )
+        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"])