Generalize CLIPArchitecture

test/architectures/test_clip.py

@@ -0,0 +1,81 @@
+# 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 pytest
+
+import torch
+from test.test_utils import assert_expected, set_rng_seed
+from torchmultimodal.architectures.clip import CLIPArchitecture
+
+
+class TestCLIPArchitecture:
+    @pytest.fixture
+    def start(self):
+        set_rng_seed(1234)
+
+        query_encoder = torch.nn.Linear(5, 3)
+        retrieval_encoder = torch.nn.Linear(4, 3)
+        encoders = torch.nn.ModuleDict(
+            {"query": query_encoder, "retrieval": retrieval_encoder}
+        )
+        clip = CLIPArchitecture(encoders=encoders)
+
+        input_query = torch.randint(1, 8, (2, 5), dtype=torch.float)
+        input_retrieval = torch.randint(1, 8, (2, 4), dtype=torch.float)
+
+        return clip, input_query, input_retrieval
+
+    def test_forward(self, start):
+        clip, input_query, input_retrieval = start
+        assert isinstance(clip, torch.nn.Module)
+
+        out = clip(modalities={"query": input_query, "retrieval": input_retrieval})
+        assert (
+            hasattr(out, "query_embeddings")
+            and hasattr(out, "retrieval_embeddings")
+            and len(out.__dict__) == 2
+        )
+
+        actual_q_embedding = out.query_embeddings
+        actual_r_embedding = out.retrieval_embeddings
+        expected_q_embedding = torch.Tensor(
+            [[-0.8066, -0.1749, 0.5647], [-0.7709, -0.1118, 0.6271]]
+        )
+        expected_r_embedding = torch.Tensor(
+            [[-0.1719, 0.7932, 0.5842], [-0.2805, 0.8761, -0.3921]]
+        )
+        assert_expected(
+            actual=actual_q_embedding, expected=expected_q_embedding, rtol=0, atol=1e-4
+        )
+        assert_expected(
+            actual=actual_r_embedding, expected=expected_r_embedding, rtol=0, atol=1e-4
+        )
+
+    def test_forward_missing_input(self, start):
+        clip, input_query, _ = start
+        assert isinstance(clip, torch.nn.Module)
+
+        with pytest.raises(AssertionError):
+            clip(modalities={"query": input_query})
+
+    def test_forward_extra_input(self, start):
+        clip, input_query, input_retrieval = start
+        assert isinstance(clip, torch.nn.Module)
+
+        with pytest.warns(UserWarning):
+            out = clip(
+                modalities={
+                    "query": input_query,
+                    "retrieval": input_retrieval,
+                    "extra": torch.Tensor([1]).to(dtype=float),
+                }
+            )
+
+        assert (
+            hasattr(out, "query_embeddings")
+            and hasattr(out, "retrieval_embeddings")
+            and len(out.__dict__) == 2
+        )

torchmultimodal/architectures/clip.py

@@ -4,6 +4,8 @@
 # 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 dataclasses import make_dataclass
 from typing import Dict
 
 import torch
@@ -20,32 +22,38 @@ class CLIPArchitecture(nn.Module):
     encoders, while the loss is implemented in ContrastiveLossWithTemperature.
 
 
-    Args:   vision_encoder (nn.Module): Instantiated vision encoder.
-                See e.g. ResNetForCLIP class.
-            text_encoder (nn.Module): Instantiated text encoder.
-                See CLIPTextEncoder class.
+    Args:   encoders (nn.ModuleDict): Dict of instantiated encoders, keyed by modality.
+                E.g. {"vision": ResNetForCLIP(), "text": CLIPTextEncoder()}
 
-    Inputs: image (Tensor): Tensor containing image features.
-            text (Tensor): Tensor containing text features.
+    Inputs: modalities (Dict[str, Tensor]): Dict of Tensor features, keyed by modality.
+                Must contain one entry for every modality in ``encoders``.
+
+    Output: CLIPOutput object with fields ``{modality}_embeddings`` for every modality
+                in ``encoders``.
     """
 
     def __init__(
         self,
-        vision_encoder: nn.Module,
-        text_encoder: nn.Module,
+        encoders: nn.ModuleDict,
     ):
         super().__init__()
-        self.vision_encoder = vision_encoder
-        self.text_encoder = text_encoder
+        self.encoders = nn.ModuleDict({k: encoders[k] for k in sorted(encoders.keys())})
 
     def forward(
         self,
-        image: torch.Tensor,
-        text: torch.Tensor,
+        modalities: Dict[str, torch.Tensor],
     ) -> Dict[str, torch.Tensor]:
-
-        img_embeddings = self.vision_encoder(image)
-        text_embeddings = self.text_encoder(text)
-        img_embeddings = F.normalize(img_embeddings)
-        text_embeddings = F.normalize(text_embeddings)
-        return {"image": img_embeddings, "text": text_embeddings}
+        embeddings = {}
+        for key, encoder in self.encoders.items():
+            assert key in modalities, f"{key} missing in input"
+            embeddings[f"{key}_embeddings"] = F.normalize(encoder(modalities[key]))
+        for key in modalities.keys():
+            if key not in self.encoders:
+                warnings.warn(f"Missing encoder for extra input {key}")
+
+        # Return a dataclass object instead of a dictionary
+        clip_output = make_dataclass(
+            "CLIPOutput",
+            [(f"{k}_embeddings", torch.Tensor) for k in self.encoders.keys()],
+        )
+        return clip_output(**embeddings)