[FLAVA]Change ordering on contrastive loss initialization

ghstack-source-id: 39675e4
Pull Request resolved: #105

Author: ankitade

test/models/flava/test_flava.py

@@ -27,27 +27,26 @@ def setUp(self):
 
     @torch.no_grad()
     def test_forward_classification(self):
-        flava = flava_model_for_classification(NUM_CLASSES)
         text = torch.randint(0, 30500, (2, 77), dtype=torch.long)
         image = torch.rand((2, 3, 224, 224))
-
         labels = torch.randint(0, 2, (2,), dtype=torch.long)
+        flava = flava_model_for_classification(NUM_CLASSES)
+        flava.eval()
 
         # Test multimodal scenario
         output = flava(image, text, "mm", labels)
-        self.assertAlmostEqual(output.loss.item(), 0.9724, places=4)
+        self.assertAlmostEqual(output.loss.item(), 0.7180, places=4)
 
         # Test unimodal image scenario
         output = flava(image, text, "image", labels)
-        self.assertAlmostEqual(output.loss.item(), 0.5453, places=4)
+        self.assertAlmostEqual(output.loss.item(), 0.7020, places=4)
 
         # Test unimodal text scenario
         output = flava(image, text, "text", labels)
-        self.assertAlmostEqual(output.loss.item(), 0.7074, places=4)
+        self.assertAlmostEqual(output.loss.item(), 0.6663, places=4)
 
     @torch.no_grad()
     def test_forward_pretraining(self):
-        flava = flava_model_for_pretraining()
         text = torch.randint(0, 30500, (2, 77), dtype=torch.long)
         image = torch.rand((2, 3, 224, 224))
         image_for_codebook = torch.rand(2, 3, 112, 112)
@@ -59,6 +58,8 @@ def test_forward_pretraining(self):
         mlm_labels[:, 1:3] = text[:, 1:3]
         itm_labels = torch.tensor((0, 1), dtype=torch.long)
 
+        flava = flava_model_for_pretraining()
+
         output = flava(
             image=image,
             text=text,
@@ -79,7 +80,7 @@ def test_forward_pretraining(self):
             sum(
                 value if value is not None else 0 for value in output.losses.values()
             ).item(),
-            20.4199,
+            21.4791,
             places=4,
         )
 
@@ -103,7 +104,7 @@ def test_forward_pretraining(self):
             sum(
                 value if value is not None else 0 for value in output.losses.values()
             ).item(),
-            9.3403,
+            8.9674,
             places=4,
         )
 
@@ -128,7 +129,7 @@ def test_forward_pretraining(self):
             sum(
                 value if value is not None else 0 for value in output.losses.values()
             ).item(),
-            10.8777,
+            10.0305,
             places=4,
         )
 

torchmultimodal/models/flava/flava_model.py

@@ -185,9 +185,8 @@ def flava_model_for_pretraining(
     # TODO: Add parameters for loss here
 ):
     model = flava_model(**flava_model_kwargs)
-
-    codebook = DalleVAEEncoder(image_size=codebook_image_size)
     losses = FLAVAPretrainingLoss()
+    codebook = DalleVAEEncoder(image_size=codebook_image_size)
 
     flava = FLAVAForPreTraining(
         model=model,
@@ -211,7 +210,7 @@ def flava_model_for_classification(
     loss_fn: Optional[Callable[..., Tensor]] = None,
     **flava_model_kwargs: Any,
 ):
-    model = flava_model(**flava_model_kwargs)
+
     classifier = MLP(
         in_dim=classifier_in_dim,
         out_dim=num_classes,
@@ -220,6 +219,7 @@ def flava_model_for_classification(
         activation=classifier_activation,
         normalization=classifier_normalization,
     )
+    model = flava_model(**flava_model_kwargs)
 
     if loss_fn is None:
         loss_fn = nn.CrossEntropyLoss()

torchmultimodal/modules/losses/flava.py

@@ -380,6 +380,21 @@ def forward(
         outputs = FLAVAPretrainingLossOutput()
         pos_mask = None
 
+        if (
+            image_sequence is not None
+            and text_sequence is not None
+            and self.contrastive_loss_weight > 0
+        ):
+            outputs.global_contrastive_output = self.contrastive_loss(
+                image_sequence,
+                text_sequence,
+                pos_mask,
+            )
+            outputs.global_contrastive_output.loss *= self.contrastive_loss_weight
+            outputs.losses.global_contrastive_loss = (
+                outputs.global_contrastive_output.loss
+            )
+
         # Check multimodal_masked_sequence to make sure this is unimodal case
         # This specific case can though be backpropagated directly as MIM is independent of
         # text, but that is a research question :)
@@ -461,19 +476,4 @@ def forward(
             outputs.mmm_image_output.loss *= self.mmm_image_loss_weight
             outputs.losses.mmm_image_loss = outputs.mmm_image_output.loss
 
-        if (
-            image_sequence is not None
-            and text_sequence is not None
-            and self.contrastive_loss_weight > 0
-        ):
-            outputs.global_contrastive_output = self.contrastive_loss(
-                image_sequence,
-                text_sequence,
-                pos_mask,
-            )
-            outputs.global_contrastive_output.loss *= self.contrastive_loss_weight
-            outputs.losses.global_contrastive_loss = (
-                outputs.global_contrastive_output.loss
-            )
-
         return outputs