Correcting cls token issues and enabling unified ViT-MAE

Summary:
1. Made cls token optional for MAE encoder and decoder.
2. Modified reconstruction lightning module to work with unified modalities
3. Added unified configs

Reviewed By: tsungyu

Differential Revision:
D49884612

Privacy Context Container: L1098566

fbshipit-source-id: 4ce66046f5002958b211778e916b3dcdd9b08755

torchmultimodal/models/masked_auto_encoder/position_embeddings.py

@@ -11,7 +11,10 @@
 
 
 def get_3d_sin_cos_embeddings(
-    embed_dim: int, temporal_size: int, spatial_size: Tuple[int, int]
+    embed_dim: int,
+    temporal_size: int,
+    spatial_size: Tuple[int, int],
+    include_cls_embed: bool = True,
 ) -> Tensor:
     """
     3d position sin cos embeddings. This implementation has been adapted from internal
@@ -20,6 +23,7 @@ def get_3d_sin_cos_embeddings(
         embed_dim (int): embedding dimension of the position embedding
         temporal_size (int): temporal input dimensions of the grid
         spatial_size (Tuple[int, int]): spatial input dimensions of the grid
+        include_cls_embed (bool): Whether to include positional embedding for [CLS] token. Defaults to True.
     return:
         embed (Tensor[int]): [1+temporal_size*spatial_size[0]*spatial_size[1], embed_dim] (w/ cls_token)
     """
@@ -60,17 +64,21 @@ def get_3d_sin_cos_embeddings(
     embed = embed.reshape([-1, embed_dim])  # [T*H*W, D]
 
     # Add pos embed for cls token
-    embed = torch.cat([torch.zeros(1, embed_dim), embed], dim=0)
+    if include_cls_embed:
+        embed = torch.cat([torch.zeros(1, embed_dim), embed], dim=0)
     embed = embed.unsqueeze(0)
     return embed
 
 
-def get_2d_sin_cos_embeddings(embed_dim: int, input_size: Tuple[int, int]) -> Tensor:
+def get_2d_sin_cos_embeddings(
+    embed_dim: int, input_size: Tuple[int, int], include_cls_embed: bool = True
+) -> Tensor:
     """
     2d position sin cos embeddings.
     Args:
         embed_dim (int): embedding dimension of the position embedding
         input_size (Tuple[int, int]): input dimensions of the grid
+        include_cls_embed (bool): Whether to include positional embedding for [CLS] token. Defaults to True.
     """
 
     # dim gets halved twice, once for h and w axis and once for sin and cos
@@ -85,7 +93,8 @@ def get_2d_sin_cos_embeddings(embed_dim: int, input_size: Tuple[int, int]) -> Te
     # h*w x embed_dim
     embed = torch.cat([embed_w, embed_h], dim=1)
     # Add pos embed for cls token
-    embed = torch.cat([torch.zeros(1, embed_dim), embed], dim=0)
+    if include_cls_embed:
+        embed = torch.cat([torch.zeros(1, embed_dim), embed], dim=0)
     embed = embed.unsqueeze(0)
     return embed