Improve support for multi-task training (#36)

Author: fcogidi

mmlearn/datasets/nyuv2.py

@@ -7,7 +7,7 @@
 import torch
 from hydra_zen import MISSING, store
 from lightning_utilities.core.imports import RequirementCache
-from PIL.Image import Image as PILImage
+from PIL import Image as PILImage
 from torch.utils.data import Dataset
 from torchvision.transforms.v2.functional import to_pil_image
 
@@ -120,7 +120,7 @@ def __init__(
 
         root_dir = os.path.join(root_dir, split)
         depth_files = [os.path.join(root_dir, "depth", f"{f}.png") for f in file_ids]
-        rgb_files = [os.path.join(root_dir, "rgb", f"{f}.jpg") for f in file_ids]
+        rgb_files = [os.path.join(root_dir, "rgb", f"{f}.png") for f in file_ids]
 
         label_files = [
             os.path.join(root_dir, "scene_class", f"{f}.txt") for f in file_ids

mmlearn/datasets/sunrgbd.py

@@ -7,7 +7,7 @@
 import torch
 from hydra_zen import MISSING, store
 from lightning_utilities.core.imports import RequirementCache
-from PIL.Image import Image as PILImage
+from PIL import Image as PILImage
 from torch.utils.data import Dataset
 from torchvision.transforms.v2.functional import to_pil_image
 

mmlearn/modules/encoders/vision.py

@@ -2,13 +2,14 @@
 
 import math
 from functools import partial
-from typing import Any, Callable, Optional, Union, cast
+from typing import Any, Callable, Literal, Optional, Union, cast
 
 import timm
 import torch
 from hydra_zen import store
 from peft import PeftConfig
 from timm.models.vision_transformer import VisionTransformer as TimmVisionTransformer
+from timm.models.vision_transformer import global_pool_nlc
 from torch import nn
 from transformers.modeling_outputs import BaseModelOutput
 
@@ -33,6 +34,9 @@ class TimmViT(nn.Module):
     ----------
     model_name : str
         The name of the model to use.
+    modality : str, default="RGB"
+        The modality of the input data. This allows this model to be used with different
+        image modalities e.g. RGB, Depth, etc.
     projection_dim : int, default=768
         The dimension of the projection head.
     pretrained : bool, default=True
@@ -51,6 +55,7 @@ class TimmViT(nn.Module):
     def __init__(
         self,
         model_name: str,
+        modality: str = "RGB",
         projection_dim: int = 768,
         pretrained: bool = True,
         freeze_layers: Union[int, float, list[int], bool] = False,
@@ -59,6 +64,7 @@ def __init__(
         model_kwargs: Optional[dict[str, Any]] = None,
     ) -> None:
         super().__init__()
+        self.modality = Modalities.get_modality(modality)
         if model_kwargs is None:
             model_kwargs = {}
 
@@ -124,7 +130,7 @@ def forward(self, inputs: dict[str, Any]) -> BaseModelOutput:
         BaseModelOutput
             The output of the model.
         """
-        x = inputs[Modalities.RGB.name]
+        x = inputs[self.modality.name]
         last_hidden_state, hidden_states = self.model.forward_intermediates(
             x, output_fmt="NLC"
         )
@@ -175,8 +181,11 @@ class VisionTransformer(nn.Module):
 
     Parameters
     ----------
-    img_size : Optional[list[int]], optional, default=None
-        list of input image sizes.
+    modality : str, optional, default="RGB"
+        The modality of the input data. This allows this model to be used with different
+        image modalities e.g. RGB, Depth, etc.
+    img_size : List[int], optional, default=None
+        List of input image sizes.
     patch_size : int, optional, default=16
         Size of each patch.
     in_chans : int, optional, default=3
@@ -209,6 +218,7 @@ class VisionTransformer(nn.Module):
 
     def __init__(
         self,
+        modality: str = "RGB",
         img_size: Optional[list[int]] = None,
         patch_size: int = 16,
         in_chans: int = 3,
@@ -218,14 +228,17 @@ def __init__(
         mlp_ratio: float = 4.0,
         qkv_bias: bool = True,
         qk_scale: Optional[float] = None,
+        global_pool: Literal["", "avg", "avgmax", "max", "token"] = "",
         drop_rate: float = 0.0,
         attn_drop_rate: float = 0.0,
         drop_path_rate: float = 0.0,
         norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
         init_std: float = 0.02,
-        **kwargs: Any,
     ) -> None:
         super().__init__()
+        assert global_pool in ("", "avg", "avgmax", "max", "token")
+
+        self.modality = Modalities.get_modality(modality)
         self.num_features = self.embed_dim = embed_dim
         self.num_heads = num_heads
         img_size = [224, 224] if img_size is None else img_size
@@ -272,6 +285,8 @@ def __init__(
         )
         self.norm = norm_layer(embed_dim)
 
+        self.global_pool = global_pool
+
         # Weight Initialization
         self.init_std = init_std
         self.apply(self._init_weights)
@@ -301,15 +316,14 @@ def _init_weights(self, m: nn.Module) -> None:
                 nn.init.constant_(m.bias, 0)
 
     def forward(
-        self,
-        x: torch.Tensor,
-        masks: Optional[Union[torch.Tensor, list[torch.Tensor]]] = None,
-        return_hidden_states: bool = False,
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, list[torch.Tensor]]]:
+        self, inputs: dict[str, Any], return_hidden_states: bool = False
+    ) -> tuple[torch.Tensor, Optional[list[torch.Tensor]]]:
         """Forward pass through the Vision Transformer."""
+        masks = inputs.get(self.modality.mask)
         if masks is not None and not isinstance(masks, list):
             masks = [masks]
 
+        x = inputs[self.modality.name]
         # -- Patchify x
         x = self.patch_embed(x)
 
@@ -336,10 +350,13 @@ def forward(
         if self.norm is not None:
             x = self.norm(x)
 
+        # -- Apply global pooling
+        x = global_pool_nlc(x, pool_type=self.global_pool)
+
         # -- Return both final output and hidden states if requested
         if return_hidden_states:
             return x, hidden_states
-        return x
+        return (x, None)
 
     def interpolate_pos_encoding(
         self, x: torch.Tensor, pos_embed: torch.Tensor
@@ -586,8 +603,7 @@ def _trunc_normal(
 def _no_grad_trunc_normal_(
     tensor: torch.Tensor, mean: float, std: float, a: float, b: float
 ) -> torch.Tensor:
-    """
-    Apply truncated normal initialization to a tensor.
+    """Apply truncated normal initialization to a tensor.
 
     Parameters
     ----------
@@ -633,15 +649,23 @@ def norm_cdf(x: float) -> float:
         provider="mmlearn",
     ),
 )
-def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
-    """
-    Create a VisionTransformerPredictor model.
+def vit_predictor(
+    kwargs: Optional[dict[str, Any]] = None,
+) -> VisionTransformerPredictor:
+    """Create a VisionTransformerPredictor model.
+
+    Parameters
+    ----------
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the predictor.
 
     Returns
     -------
     VisionTransformerPredictor
         An instance of VisionTransformerPredictor.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformerPredictor(
         mlp_ratio=4, qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs
     )
@@ -654,15 +678,25 @@ def vit_predictor(**kwargs: Any) -> VisionTransformerPredictor:
         provider="mmlearn",
     ),
 )
-def vit_tiny(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with tiny configuration.
+def vit_tiny(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with tiny configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=192,
@@ -682,15 +716,25 @@ def vit_tiny(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
         provider="mmlearn",
     ),
 )
-def vit_small(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with small configuration.
+def vit_small(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with small configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=384,
@@ -710,15 +754,25 @@ def vit_small(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
         provider="mmlearn",
     ),
 )
-def vit_base(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with base configuration.
+def vit_base(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with base configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=768,
@@ -738,15 +792,25 @@ def vit_base(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
         provider="mmlearn",
     ),
 )
-def vit_large(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with large configuration.
+def vit_large(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with large configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=1024,
@@ -766,15 +830,25 @@ def vit_large(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
         provider="mmlearn",
     ),
 )
-def vit_huge(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with huge configuration.
+def vit_huge(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with huge configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=1280,
@@ -794,15 +868,25 @@ def vit_huge(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
         provider="mmlearn",
     ),
 )
-def vit_giant(patch_size: int = 16, **kwargs: Any) -> VisionTransformer:
-    """
-    Create a VisionTransformer model with giant configuration.
+def vit_giant(
+    patch_size: int = 16, kwargs: Optional[dict[str, Any]] = None
+) -> VisionTransformer:
+    """Create a VisionTransformer model with giant configuration.
+
+    Parameters
+    ----------
+    patch_size : int, default=16
+        Size of each patch.
+    kwargs : dict[str, Any], optional, default=None
+        Keyword arguments for the model variant.
 
     Returns
     -------
     VisionTransformer
         An instance of VisionTransformer.
     """
+    if kwargs is None:
+        kwargs = {}
     return VisionTransformer(
         patch_size=patch_size,
         embed_dim=1408,