[feat] Add VinVL model wrapper (#1151)

Summary:
Pull Request resolved: #1151

Add VinVL 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 BaseModel instantiation from config,
and forward pass for classification and pretraining.

Reviewed By: ebsmothers

Differential Revision: D32574738

Pulled By: Ryan-Qiyu-Jiang

fbshipit-source-id: 50f8396821effd778c6d5184cd940864fc1eb3b1

Author: Ryan-Qiyu-Jiang

mmf/configs/models/vinvl/defaults.yaml

@@ -0,0 +1,18 @@
+model_config:
+  vinvl:
+    heads:
+      test:
+        type: mlp
+        freeze: false
+        lr_multiplier: 1.0
+        in_dim: 768
+        hidden_size: 1536
+        num_labels: 3129
+        pooler_name: bert_pooler
+    bert_model_name: bert-base-uncased
+    loss_type: sfmx
+    img_feature_dim: 2054
+    img_feature_type: 'frcnn'
+    use_img_layernorm: 1
+    img_layer_norm_eps: 1e-12
+    max_img_seq_len: 70

mmf/models/__init__.py

@@ -18,6 +18,7 @@
 from .uniter import UNITER
 from .vilbert import ViLBERT
 from .vilt import ViLT
+from .vinvl import VinVL
 from .visual_bert import VisualBERT
 
 __all__ = [
@@ -46,4 +47,5 @@
     "AlbefVitEncoder",
     "ViLT",
     "UNITER",
+    "VinVL",
 ]

mmf/models/vinvl.py

@@ -6,15 +6,18 @@
 
 import logging
 from collections import namedtuple
-from dataclasses import asdict
-from typing import Dict, Optional, Tuple
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Optional, Tuple
 
 import torch
+from mmf.common.registry import registry
 from mmf.common.sample import SampleList
+from mmf.models.base_model import BaseModel
 from mmf.models.transformers.heads.contrastive import ThreeWayContrastive
 from mmf.models.transformers.heads.mlm import MLM
 from mmf.models.transformers.heads.mlp import MLP
 from mmf.utils.general import retry_n
+from omegaconf import MISSING, OmegaConf
 from torch import Tensor, nn
 from transformers.modeling_bert import (
     BertConfig,
@@ -378,3 +381,119 @@ def forward(
         )
         losses = {**mlm_result, **contrastive_loss_result}
         return {"losses": losses}
+
+
+@registry.register_model("vinvl")
+class VinVL(BaseModel):
+    """VinVL base model called by MMF.
+    VinVL paper, 3-way contrastive loss:
+    https://arxiv.org/pdf/2101.00529.pdf
+
+    Implementation based on https://github.com/microsoft/Oscar
+
+    Expects VinVL features extracted by
+    https://github.com/microsoft/scene_graph_benchmark
+    using Visual Genome object detection labels.
+
+    The label map used for training is available at
+    https://github.com/microsoft/scene_graph_benchmark/blob/main/README.md
+    """
+
+    @dataclass
+    class Config:
+        random_init: bool = False
+        bert_model_name: str = "bert-base-uncased"
+        hidden_size: int = 768
+        heads: Any = MISSING
+        do_pretraining: bool = False
+        img_feature_dim: int = 2054
+        img_feature_type: str = "frcnn"
+        use_img_layernorm: bool = True
+        img_layer_norm_eps: float = 1e-12
+        max_img_seq_len: int = 70
+
+    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/vinvl/defaults.yaml"
+
+    def build(self):
+        if self.do_pretraining:
+            mlm_config = self.config.heads.get("mlm")
+            contrast_config = self.config.heads.get("contrast")
+            self.vinvl = VinVLForPretraining(
+                mlm_config=mlm_config, contrast_config=contrast_config, **self.config
+            )
+        else:
+            # do classification finetuning
+            mlp_config = self.config.heads.get("mlp")
+            loss_config = self.config.get("ce_loss")
+            self.vinvl = VinVLForClassification(
+                mlp_config=mlp_config, loss_config=loss_config, **self.config
+            )
+
+    def init_losses(self):
+        """
+        Defer loss management to submodels,
+        do nothing when called by build_model.
+        """
+
+    def forward(self, sample_list: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        attention_mask = self._get_attention_mask(
+            sample_list["image_feature_0"],
+            sample_list["image_info_0"],
+            sample_list["input_mask"],
+        )
+
+        if self.do_pretraining:
+            corrupt_attention_mask = self._get_attention_mask(
+                sample_list["image_feature_0"],
+                sample_list["image_info_0"],
+                sample_list["input_mask_corrupt"],
+            )
+            return self.vinvl(
+                sample_list["input_ids_masked"],
+                sample_list["input_ids_corrupt"],
+                sample_list["lm_label_ids"],
+                sample_list["contrastive_labels"],
+                sample_list["segment_ids"],
+                attention_mask,
+                sample_list["segment_ids_corrupt"],
+                corrupt_attention_mask,
+                sample_list["image_feature_0"],
+            )
+        else:
+            return self.vinvl(
+                sample_list["input_ids"],
+                sample_list["segment_ids"],
+                attention_mask,
+                sample_list["image_feature_0"],
+                labels=sample_list.get("labels"),
+            )
+
+    def _get_attention_mask(
+        self, image_feat: Tensor, image_info: Dict[str, Tensor], input_mask: Tensor
+    ) -> Tensor:
+        # image_dim = (bs,)
+        # with the number of features per image in the batch as an int
+        image_dim = image_info.get("max_features")
+        if image_dim is None:
+            image_mask = torch.ones(
+                (image_feat.size(0), image_feat.size(1)), device=image_feat.device
+            ).long()
+        else:
+            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
+            image_mask = image_mask.long()
+
+        attention_mask = torch.cat((input_mask, image_mask), dim=-1)
+        return attention_mask

tests/models/test_vinvl.py

@@ -2,13 +2,19 @@
 
 import unittest
 
+import tests.test_utils as test_utils
 import torch
+from mmf.common.sample import SampleList
 from mmf.models.vinvl import (
     VinVLBase,
     VinVLForClassification,
     VinVLForPretraining,
 )
+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
 from transformers.modeling_bert import BertConfig
 
 
@@ -93,3 +99,99 @@ def test_pretraining_forward(self):
         self.assertTrue("losses" in model_output)
         self.assertTrue("masked_lm_loss" in model_output["losses"])
         self.assertTrue("three_way_contrastive_loss" in model_output["losses"])
+
+
+class TestVinVLModel(unittest.TestCase):
+    def setUp(self):
+        test_utils.setup_proxy()
+        setup_imports()
+        model_name = "vinvl"
+        args = test_utils.dummy_args(model=model_name, dataset="test")
+        configuration = Configuration(args)
+        config = configuration.get_config()
+        model_config = config.model_config[model_name]
+        model_config.model = model_name
+        model_config.do_pretraining = False
+        classification_config_dict = {
+            "do_pretraining": False,
+            "heads": {"mlp": {"num_labels": 3129}},
+            "ce_loss": {"ignore_index": -1},
+        }
+        self.classification_config = OmegaConf.create(
+            {**model_config, **classification_config_dict}
+        )
+
+        pretraining_config_dict = {
+            "do_pretraining": True,
+            "heads": {"mlm": {"hidden_size": 768}},
+        }
+        self.pretraining_config = OmegaConf.create(
+            {**model_config, **pretraining_config_dict}
+        )
+
+        self.sample_list = self._get_sample_list()
+
+    def tearDown(self):
+        teardown_imports()
+
+    def _get_sample_list(self):
+        bs = 8
+        num_feats = 70
+
+        class MockObj:
+            pass
+
+        mock_input = MockObj()
+        mock_vinvl_input_tensors(mock_input, bs=bs, num_feats=num_feats)
+
+        input_mask = torch.ones_like(mock_input.input_ids)
+        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,
+        }
+
+        sample_list = SampleList()
+        sample_list.add_field("input_ids", mock_input.input_ids)
+        sample_list.add_field("input_ids_corrupt", mock_input.input_ids)
+        sample_list.add_field("input_ids_masked", mock_input.input_ids)
+        sample_list.add_field("image_feature_0", mock_input.img_feats)
+        sample_list.add_field("image_info_0", image_info)
+        sample_list.add_field("input_mask", input_mask)
+        sample_list.add_field("input_mask_corrupt", input_mask)
+        sample_list.add_field("segment_ids", mock_input.token_type_ids)
+        sample_list.add_field("segment_ids_corrupt", mock_input.token_type_ids)
+        sample_list.add_field("labels", mock_input.labels)
+        sample_list.add_field("contrastive_labels", mock_input.contrastive_labels)
+        sample_list.add_field("lm_label_ids", mock_input.lm_label_ids)
+        sample_list = sample_list.to(get_current_device())
+        sample_list.dataset_name = "test"
+        sample_list.dataset_type = "test"
+        return sample_list
+
+    def test_vinvl_for_classification(self):
+        model_for_classification = build_model(self.classification_config)
+        model_for_classification.eval()
+        model_for_classification = model_for_classification.to(get_current_device())
+        with torch.no_grad():
+            model_output = model_for_classification(self.sample_list)
+
+        self.assertTrue("losses" in model_output)
+        self.assertTrue("ce" in model_output["losses"])
+
+    def test_vinvl_for_pretraining(self):
+        model_for_pretraining = build_model(self.pretraining_config)
+        model_for_pretraining.eval()
+        model_for_pretraining = model_for_pretraining.to(get_current_device())
+
+        with torch.no_grad():
+            model_output = model_for_pretraining(self.sample_list)
+
+        self.assertTrue("losses" in model_output)
+        self.assertTrue("masked_lm_loss" in model_output["losses"])
+        self.assertTrue("three_way_contrastive_loss" in model_output["losses"])