[FLAVA]Separate the pretraining loss from the pretraininig model

ghstack-source-id: 417f0746292f666c28f5092eb4545ab0dca567c0
Pull Request resolved: #278

examples/flava/model.py

@@ -8,16 +8,18 @@
 
 import torch
 from pytorch_lightning import LightningModule
+from torch import nn
 from torchmetrics import Accuracy
 from torchmultimodal.models.flava.model import (
  flava_model_for_classification,
  flava_model_for_pretraining,
 )
+from torchmultimodal.modules.losses.flava import FLAVAPretrainingLoss
 from transformers.optimization import get_cosine_schedule_with_warmup
 
 
 def get_optimizers_for_lightning(
- model: torch.nn.Module,
+ model: nn.Module,
  learning_rate: float,
  adam_eps: float,
  adam_weight_decay: float,
@@ -59,6 +61,7 @@ def __init__(
  self.adam_weight_decay = adam_weight_decay
  self.warmup_steps = warmup_steps
  self.max_steps = max_steps
+ self.loss = FLAVAPretrainingLoss(logit_scale=self.model.logit_scale)
 
  def training_step(self, batch, batch_idx):
  output = self._step(batch, batch_idx)
@@ -104,7 +107,24 @@ def _step(self, batch, batch_idx):
  itm_labels=batch.get("itm_labels", None),
  required_embedding=required_embedding,
  )
- return output
+
+ loss = self.loss(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
+ return loss
 
  def configure_optimizers(self):
  return get_optimizers_for_lightning(

test/models/flava/test_checkpoint.py

@@ -12,6 +12,7 @@
  flava_model_for_classification,
  flava_model_for_pretraining,
 )
+from torchmultimodal.modules.losses.flava import FLAVAPretrainingLoss
 
 
 @pytest.fixture(autouse=True)
@@ -139,8 +140,25 @@ def test_flava_model_for_pretraining(self, inputs_pretraining, pretraining_model
  image_input = inputs_pretraining("image")
  text_input = inputs_pretraining("text")
  flava = pretraining_model()
-
+ losses = FLAVAPretrainingLoss(flava.logit_scale)
  output = flava(*mm_input)
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
+
  actual = output.losses
  expected = dict(
  mmm_text_loss=10.9567,
@@ -153,6 +171,22 @@ def test_flava_model_for_pretraining(self, inputs_pretraining, pretraining_model
  self._assert_tensor_dicts_equal(actual, expected)
 
  output = flava(*image_input)
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
  actual = output.losses
  expected = dict(
  mmm_text_loss=None,
@@ -165,6 +199,22 @@ def test_flava_model_for_pretraining(self, inputs_pretraining, pretraining_model
  self._assert_tensor_dicts_equal(actual, expected)
 
  output = flava(*text_input)
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
  actual = output.losses
  expected = dict(
  mmm_text_loss=None,

test/models/flava/test_flava.py

@@ -18,6 +18,7 @@
  FLAVAOutput,
 )
 from torchmultimodal.modules.layers.transformer import TransformerOutput
+from torchmultimodal.modules.losses.flava import FLAVAPretrainingLoss
 
 NUM_CLASSES = 2
 
@@ -61,6 +62,7 @@ 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()
+ losses = FLAVAPretrainingLoss(flava.logit_scale)
  flava.eval()
  output = flava(
  image=image,
@@ -72,6 +74,22 @@ def test_forward_pretraining(self):
  itm_labels=itm_labels,
  mlm_labels=mlm_labels,
  )
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
  self.assertIsNone(output.mlm_output)
  self.assertIsNone(output.mim_output)
  self.assertIsNotNone(output.global_contrastive_output)
@@ -96,6 +114,22 @@ def test_forward_pretraining(self):
  itm_labels=itm_labels,
  mlm_labels=mlm_labels,
  )
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
  self.assertIsNone(output.mlm_output)
  self.assertIsNotNone(output.mim_output)
  self.assertIsNone(output.global_contrastive_output)
@@ -120,6 +154,22 @@ def test_forward_pretraining(self):
  itm_labels=itm_labels,
  mlm_labels=mlm_labels,
  )
+ output = losses(
+ multimodal_masked_sequence=output.multimodal_masked_sequence,
+ pos_mask=output.pos_mask,
+ itm_labels=output.itm_labels,
+ mim_labels=output.mim_labels,
+ mlm_labels=output.mlm_labels,
+ mmm_mlm_labels=output.mmm_mlm_labels,
+ mmm_mim_labels=output.mmm_mim_labels,
+ projected_image_embeddings=output.projected_image_embeddings,
+ projected_text_embeddings=output.projected_text_embeddings,
+ itm_logits=output.itm_logits,
+ mlm_head_output=output.mlm_head_output,
+ mim_head_output=output.mim_head_output,
+ mmm_mlm_head_output=output.mmm_mlm_head_output,
+ mmm_mim_head_output=output.mmm_mim_head_output,
+ )
  self.assertIsNotNone(output.mlm_output)
  self.assertIsNone(output.mim_output)
  self.assertIsNone(output.global_contrastive_output)

torchmultimodal/models/flava/model.py

@@ -24,11 +24,7 @@
  TransformerEncoder,
  TransformerOutput,
 )
-from torchmultimodal.modules.losses.flava import (
- FLAVAPretrainingLoss,
- FLAVAPretrainingLossOutput,
- Pooler,
-)
+from torchmultimodal.modules.losses.flava import FLAVAPretrainingLossOutput, Pooler
 from torchmultimodal.utils.common import ModelOutput, PretrainedMixin
 from typing_extensions import Literal
 
@@ -62,7 +58,7 @@
 FLAVA_FOR_PRETRAINED_MAPPING = {
  # This will no longer load with the updated model, but keeping here just in case
  # "flava_full": "https://huggingface.co/aps/flava_full_pretrained_encoders_torchmm/resolve/main/pytorch_model.bin",
- "flava_full": "https://download.pytorch.org/models/multimodal/flava/flava_for_pretraining_unified_itm_mp.pt",
+ "flava_full": "https://download.pytorch.org/models/multimodal/flava/flava_for_pretraining_no_loss.pt",
 }
 
 FLAVA_MODEL_MAPPING = {
@@ -105,6 +101,24 @@ class FLAVAForClassificationOutput(ModelOutput):
  loss: Tensor
 
 
+@dataclass
+class FLAVAForPretrainingOutput:
+ multimodal_masked_sequence: Tensor
+ pos_mask: Tensor
+ mim_labels: Tensor
+ mlm_labels: Tensor
+ mmm_mlm_labels: Tensor
+ mmm_mim_labels: Tensor
+ itm_labels: Tensor
+ projected_image_embeddings: Tensor
+ projected_text_embeddings: Tensor
+ itm_logits: Tensor
+ mlm_head_output: Tensor
+ mim_head_output: Tensor
+ mmm_mlm_head_output: Tensor
+ mmm_mim_head_output: Tensor
+
+
 class FLAVAModel(nn.Module, PretrainedMixin):
  def __init__(
  self,
@@ -367,22 +381,22 @@ def __init__(
  self,
  model: FLAVAModel,
  image_codebook: nn.Module,
- loss: nn.Module,
  itm_head: nn.Module,
  mlm_head: nn.Module,
  mim_head: nn.Module,
  mmm_mlm_head: nn.Module,
  mmm_mim_head: nn.Module,
+ logit_scale: nn.Module,
  ):
  super().__init__()
  self.model = model
  self.image_codebook = image_codebook
- self.loss = loss
  self.itm_head = itm_head
  self.mlm_head = mlm_head
  self.mim_head = mim_head
  self.mmm_mlm_head = mmm_mlm_head
  self.mmm_mim_head = mmm_mim_head
+ self.logit_scale = logit_scale
 
  def encode_image(
  self,
@@ -469,8 +483,10 @@ def forward(
  itm_logits = self.itm_head(multimodal_masked_sequence)
 
  multimodal_masked_sequence = multimodal_masked_sequence[pos_mask]
+
  if mlm_labels is not None:
  mmm_mlm_labels = mlm_labels[pos_mask]
+
  if image_labels is not None:
  mmm_mim_labels = image_labels[pos_mask]
 
@@ -494,14 +510,14 @@ def forward(
  sequence_for_image = multimodal_masked_sequence[:, 2 : 2 + total_indices, :]
  mmm_mim_head_output = self.mmm_mim_head(sequence_for_image, mmm_mim_labels)
 
- return self.loss(
+ return FLAVAForPretrainingOutput(
  multimodal_masked_sequence=flava_output.multimodal_masked.last_hidden_state,
  pos_mask=pos_mask,
- itm_labels=itm_labels,
  mim_labels=image_labels,
  mlm_labels=mlm_labels,
- mmm_mlm_labels=mmm_mlm_labels,
  mmm_mim_labels=mmm_mim_labels,
+ mmm_mlm_labels=mmm_mlm_labels,
+ itm_labels=itm_labels,
  projected_image_embeddings=flava_output.projected_image_embeddings,
  projected_text_embeddings=flava_output.projected_text_embeddings,
  itm_logits=itm_logits,
@@ -660,6 +676,7 @@ def flava_model_for_pretraining(
  codebook_image_size: int = 112,
  pretrained_model_key: Optional[str] = None,
  image_vocab_size: int = 8192,
+ logit_scale: float = math.log(1 / 0.07),
  **flava_model_kwargs: Any,
  # TODO: Add parameters for loss here
 ) -> FLAVAForPreTraining:
@@ -677,18 +694,18 @@ def flava_model_for_pretraining(
  mmm_mim_head = MaskedPredictionHead(
  hidden_size=hidden_size, vocab_size=image_vocab_size
  )
- losses = FLAVAPretrainingLoss()
+
  codebook = DalleVAEEncoder(image_size=codebook_image_size)
 
  flava = FLAVAForPreTraining(
  model=model,
  image_codebook=codebook,
- loss=losses,
  itm_head=itm_head,
  mlm_head=mlm_head,
  mim_head=mim_head,
  mmm_mlm_head=mmm_mlm_head,
  mmm_mim_head=mmm_mim_head,
+ logit_scale=nn.Parameter(logit_scale * torch.ones([])),
  )
 
  if pretrained_model_key is not None: