Remove ContrastiveLossBase for just a function

Author: bdvllrs

shimmer/__init__.py

@@ -1,8 +1,9 @@
 from shimmer.config import (ShimmerInfoConfig, load_config,
                             load_structured_config)
 from shimmer.modules.contrastive_loss import (
-    ContrastiveLoss, ContrastiveLossBase, ContrastiveLossWithUncertainty,
-    contrastive_loss, contrastive_loss_with_uncertainty)
+    ContrastiveLoss, ContrastiveLossType, ContrastiveLossWithUncertainty,
+    VarContrastiveLossType, contrastive_loss,
+    contrastive_loss_with_uncertainty)
 from shimmer.modules.domain import DomainModule, LossOutput
 from shimmer.modules.global_workspace import (GlobalWorkspace,
                                               GlobalWorkspaceBase,
@@ -39,7 +40,8 @@
     "VariationalGWInterface",
     "VariationalGWModule",
     "ContrastiveLoss",
-    "ContrastiveLossBase",
+    "ContrastiveLossType",
+    "VarContrastiveLossType",
     "ContrastiveLossWithUncertainty",
     "contrastive_loss",
     "contrastive_loss_with_uncertainty",

shimmer/modules/__init__.py

@@ -1,6 +1,7 @@
 from shimmer.modules.contrastive_loss import (
-    ContrastiveLoss, ContrastiveLossBase, ContrastiveLossWithUncertainty,
-    contrastive_loss, contrastive_loss_with_uncertainty)
+    ContrastiveLoss, ContrastiveLossType, ContrastiveLossWithUncertainty,
+    VarContrastiveLossType, contrastive_loss,
+    contrastive_loss_with_uncertainty)
 from shimmer.modules.domain import DomainModule, LossOutput
 from shimmer.modules.global_workspace import (GlobalWorkspace,
                                               GlobalWorkspaceBase,
@@ -30,11 +31,12 @@
     "VariationalGWInterface",
     "VariationalGWModule",
     "ContrastiveLoss",
-    "ContrastiveLossBase",
+    "ContrastiveLossType",
+    "VarContrastiveLossType",
     "ContrastiveLossWithUncertainty",
     "contrastive_loss",
     "contrastive_loss_with_uncertainty",
-    "LatentT",
+    "LatentsT",
     "LatentsDomainGroupT",
     "GWLosses",
     "GWLossesBase",

shimmer/modules/contrastive_loss.py

@@ -1,12 +1,16 @@
-from abc import ABC, abstractmethod
-from collections.abc import Mapping
-from typing import Literal, TypedDict
+from collections.abc import Callable
+from typing import Literal
 
 import torch
 from torch.nn.functional import cross_entropy, normalize
 
 from shimmer.modules.domain import LossOutput
 
+ContrastiveLossType = Callable[[torch.Tensor, torch.Tensor], LossOutput]
+VarContrastiveLossType = Callable[
+    [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], LossOutput
+]
+
 
 def info_nce(
     x: torch.Tensor,
@@ -58,13 +62,7 @@ def contrastive_loss_with_uncertainty(
     return 0.5 * (ce + ce_t)
 
 
-class ContrastiveLossBase(torch.nn.Module, ABC):
-    @abstractmethod
-    def forward(self, x: torch.Tensor, y: torch.Tensor) -> LossOutput:
-        ...
-
-
-class ContrastiveLoss(ContrastiveLossBase):
+class ContrastiveLoss(torch.nn.Module):
     logit_scale: torch.Tensor
 
     def __init__(
@@ -77,17 +75,14 @@ def __init__(
         self.register_buffer("logit_scale", logit_scale)
         self.reduction: Literal["mean", "sum", "none"] = reduction
 
-    def __call__(self, *args, **kwargs) -> LossOutput:
-        return super().__call__(*args, **kwargs)
-
     def forward(self, x: torch.Tensor, y: torch.Tensor) -> LossOutput:
         return LossOutput(
             contrastive_loss(x, y, self.logit_scale, self.reduction),
             {"logit_scale": self.logit_scale.exp()},
         )
 
 
-class ContrastiveLossWithUncertainty(ContrastiveLossBase):
+class ContrastiveLossWithUncertainty(torch.nn.Module):
     logit_scale: torch.Tensor
 
     def __init__(

shimmer/modules/global_workspace.py

@@ -9,7 +9,7 @@
 from torch.optim.lr_scheduler import OneCycleLR
 
 from shimmer.modules.contrastive_loss import (ContrastiveLoss,
-                                              ContrastiveLossBase,
+                                              ContrastiveLossType,
                                               ContrastiveLossWithUncertainty)
 from shimmer.modules.dict_buffer import DictBuffer
 from shimmer.modules.domain import DomainModule
@@ -316,18 +316,24 @@ def __init__(
         domain_mods = freeze_domain_modules(domain_mods)
         coef_buffers = DictBuffer(loss_coefs)
 
-        contrastive_fn_class = (
-            ContrastiveLossWithUncertainty
-            if var_contrastive_loss
-            else ContrastiveLoss
-        )
-
-        loss_mod = VariationalGWLosses(
-            gw_mod,
-            domain_mods,
-            coef_buffers,
-            contrastive_fn_class(torch.tensor([1]).log(), "mean"),
-        )
+        if var_contrastive_loss:
+            loss_mod = VariationalGWLosses(
+                gw_mod,
+                domain_mods,
+                coef_buffers,
+                var_contrastive_fn=ContrastiveLossWithUncertainty(
+                    torch.tensor([1]).log(), "mean"
+                ),
+            )
+        else:
+            loss_mod = VariationalGWLosses(
+                gw_mod,
+                domain_mods,
+                coef_buffers,
+                contrastive_fn=ContrastiveLoss(
+                    torch.tensor([1]).log(), "mean"
+                ),
+            )
 
         super().__init__(
             gw_mod,
@@ -346,14 +352,14 @@ def pretrained_global_workspace(
     gw_interfaces: Mapping[str, GWInterfaceBase],
     workspace_dim: int,
     loss_coefs: Mapping[str, torch.Tensor],
-    var_contrastive_loss: bool = False,
+    contrastive_fn: ContrastiveLossType,
     **kwargs,
 ) -> GlobalWorkspaceBase:
     gw_mod = VariationalGWModule(gw_interfaces, workspace_dim)
     domain_mods = freeze_domain_modules(domain_mods)
     coef_buffers = DictBuffer(loss_coefs)
     loss_mod = VariationalGWLosses(
-        gw_mod, domain_mods, coef_buffers, var_contrastive_loss
+        gw_mod, domain_mods, coef_buffers, contrastive_fn
     )
 
     return cast(

shimmer/modules/losses.py

@@ -3,7 +3,8 @@
 
 import torch
 
-from shimmer.modules.contrastive_loss import ContrastiveLossBase
+from shimmer.modules.contrastive_loss import (ContrastiveLossType,
+                                              VarContrastiveLossType)
 from shimmer.modules.dict_buffer import DictBuffer
 from shimmer.modules.domain import DomainModule
 from shimmer.modules.gw_module import (GWModule, GWModuleBase,
@@ -156,7 +157,7 @@ def _translation_loss(
 def _contrastive_loss(
     gw_mod: GWModuleBase,
     latent_domains: LatentsT,
-    contrastive_fn: ContrastiveLossBase,
+    contrastive_fn: ContrastiveLossType,
 ) -> dict[str, torch.Tensor]:
     losses: dict[str, torch.Tensor] = {}
     metrics: dict[str, torch.Tensor] = {}
@@ -197,7 +198,7 @@ def _contrastive_loss(
 def _contrastive_loss_with_uncertainty(
     gw_mod: VariationalGWModule,
     latent_domains: LatentsT,
-    contrastive_fn: ContrastiveLossBase,
+    contrastive_fn: VarContrastiveLossType,
 ) -> dict[str, torch.Tensor]:
     losses: dict[str, torch.Tensor] = {}
     metrics: dict[str, torch.Tensor] = {}
@@ -246,7 +247,7 @@ def __init__(
         gw_mod: GWModule,
         domain_mods: dict[str, DomainModule],
         coef_buffers: DictBuffer,
-        contrastive_fn: ContrastiveLossBase,
+        contrastive_fn: ContrastiveLossType,
     ):
         super().__init__()
         self.gw_mod = gw_mod
@@ -303,14 +304,19 @@ def __init__(
         gw_mod: VariationalGWModule,
         domain_mods: dict[str, DomainModule],
         coef_buffers: DictBuffer,
-        contrastive_fn: ContrastiveLossBase,
+        contrastive_fn: ContrastiveLossType | None = None,
+        var_contrastive_fn: VarContrastiveLossType | None = None,
     ):
         super().__init__()
 
         self.gw_mod = gw_mod
         self.domain_mods = domain_mods
         self.loss_coefs = coef_buffers
+        assert (contrastive_fn is not None) != (
+            var_contrastive_fn is not None
+        ), "Should either have contrastive_fn or var_contrastive_fn"
         self.contrastive_fn = contrastive_fn
+        self.var_contrastive_fn = var_contrastive_fn
 
     def demi_cycle_loss(
         self, latent_domains: LatentsT
@@ -328,10 +334,12 @@ def translation_loss(
     def contrastive_loss(
         self, latent_domains: LatentsT
     ) -> dict[str, torch.Tensor]:
-        if self.var_contrastive_loss:
+        if self.var_contrastive_fn is not None:
             return _contrastive_loss_with_uncertainty(
-                self.gw_mod, latent_domains, self.contrastive_fn
+                self.gw_mod, latent_domains, self.var_contrastive_fn
             )
+
+        assert self.contrastive_fn is not None
         return _contrastive_loss(
             self.gw_mod, latent_domains, self.contrastive_fn
         )