Add pretrained_global_workspace

ruflab · Jan 22, 2024 · 7161b16 · 7161b16
1 parent ac0cd34
commit 7161b16
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 16 deletions.
diff --git a/shimmer/__init__.py b/shimmer/__init__.py
@@ -3,9 +3,11 @@
 from shimmer.modules.domain import DomainModule
 from shimmer.modules.global_workspace import (GlobalWorkspace,
                                               GlobalWorkspaceBase,
-                                              VariationalGlobalWorkspace)
-from shimmer.modules.gw_module import (BaseGWInterface, DeterministicGWModule,
-                                       GWDecoder, GWEncoder, GWInterface,
+                                              SchedulerArgs,
+                                              VariationalGlobalWorkspace,
+                                              pretrained_global_workspace)
+from shimmer.modules.gw_module import (DeterministicGWModule, GWDecoder,
+                                       GWEncoder, GWInterface, GWInterfaceBase,
                                        GWModule, VariationalGWEncoder,
                                        VariationalGWInterface,
                                        VariationalGWModule)
@@ -19,7 +21,7 @@
     "load_structured_config",
     "ShimmerInfoConfig",
     "DomainModule",
-    "BaseGWInterface",
+    "GWInterfaceBase",
     "DeterministicGWModule",
     "GWDecoder",
     "GWEncoder",
@@ -34,4 +36,6 @@
     "GlobalWorkspace",
     "GlobalWorkspaceBase",
     "VariationalGlobalWorkspace",
+    "SchedulerArgs",
+    "pretrained_global_workspace",
 ]
diff --git a/shimmer/modules/__init__.py b/shimmer/modules/__init__.py
@@ -1,9 +1,11 @@
 from shimmer.modules.domain import DomainModule
 from shimmer.modules.global_workspace import (GlobalWorkspace,
                                               GlobalWorkspaceBase,
-                                              VariationalGlobalWorkspace)
-from shimmer.modules.gw_module import (BaseGWInterface, DeterministicGWModule,
-                                       GWDecoder, GWEncoder, GWInterface,
+                                              SchedulerArgs,
+                                              VariationalGlobalWorkspace,
+                                              pretrained_global_workspace)
+from shimmer.modules.gw_module import (DeterministicGWModule, GWDecoder,
+                                       GWEncoder, GWInterface, GWInterfaceBase,
                                        GWModule, VariationalGWEncoder,
                                        VariationalGWInterface,
                                        VariationalGWModule)
@@ -12,7 +14,7 @@
 
 __all__ = [
     "DomainModule",
-    "BaseGWInterface",
+    "GWInterfaceBase",
     "DeterministicGWModule",
     "GWDecoder",
     "GWEncoder",
@@ -27,4 +29,6 @@
     "GlobalWorkspace",
     "GlobalWorkspaceBase",
     "VariationalGlobalWorkspace",
+    "SchedulerArgs",
+    "pretrained_global_workspace",
 ]
diff --git a/shimmer/modules/global_workspace.py b/shimmer/modules/global_workspace.py
@@ -1,4 +1,5 @@
 from collections.abc import Iterable, Mapping
+from pathlib import Path
 from typing import Any, TypedDict, cast
 
 import torch
@@ -9,7 +10,7 @@
 
 from shimmer.modules.dict_buffer import DictBuffer
 from shimmer.modules.domain import DomainModule
-from shimmer.modules.gw_module import (BaseGWInterface, DeterministicGWModule,
+from shimmer.modules.gw_module import (DeterministicGWModule, GWInterfaceBase,
                                        GWModule, VariationalGWModule)
 from shimmer.modules.losses import (DeterministicGWLosses, GWLosses, LatentsT,
                                     VariationalGWLosses)
@@ -268,7 +269,7 @@ class GlobalWorkspace(GlobalWorkspaceBase):
     def __init__(
         self,
         domain_mods: Mapping[str, DomainModule],
-        gw_interfaces: Mapping[str, BaseGWInterface],
+        gw_interfaces: Mapping[str, GWInterfaceBase],
         workspace_dim: int,
         loss_coefs: dict[str, torch.Tensor],
         optim_lr: float = 1e-3,
@@ -295,7 +296,7 @@ class VariationalGlobalWorkspace(GlobalWorkspaceBase):
     def __init__(
         self,
         domain_mods: Mapping[str, DomainModule],
-        gw_interfaces: Mapping[str, BaseGWInterface],
+        gw_interfaces: Mapping[str, GWInterfaceBase],
         workspace_dim: int,
         loss_coefs: dict[str, torch.Tensor],
         var_contrastive_loss: bool = False,
@@ -319,3 +320,32 @@ def __init__(
             optim_weight_decay,
             scheduler_args,
         )
+
+
+def pretrained_global_workspace(
+    checkpoint_path: str | Path,
+    domain_mods: Mapping[str, DomainModule],
+    gw_interfaces: Mapping[str, GWInterfaceBase],
+    workspace_dim: int,
+    loss_coefs: dict[str, torch.Tensor],
+    var_contrastive_loss: bool = False,
+    **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
+    )
+
+    return cast(
+        GlobalWorkspaceBase,
+        GlobalWorkspaceBase.load_from_checkpoint(
+            checkpoint_path,
+            domain_mods=domain_mods,
+            gw_mod=gw_mod,
+            coef_buffers=coef_buffers,
+            loss_mod=loss_mod,
+            **kwargs,
+        ),
+    )
diff --git a/shimmer/modules/gw_module.py b/shimmer/modules/gw_module.py
@@ -80,7 +80,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         return self.layers(x), self.uncertainty_level.expand(x.size(0), -1)
 
 
-class BaseGWInterface(nn.Module, ABC):
+class GWInterfaceBase(nn.Module, ABC):
     def __init__(
         self, domain_module: DomainModule, workspace_dim: int
     ) -> None:
@@ -99,12 +99,12 @@ def decode(self, z: torch.Tensor) -> torch.Tensor:
 
 class GWModule(nn.Module, ABC):
     def __init__(
-        self, gw_interfaces: Mapping[str, BaseGWInterface], workspace_dim: int
+        self, gw_interfaces: Mapping[str, GWInterfaceBase], workspace_dim: int
     ) -> None:
         super().__init__()
         # casting for LSP autocompletion
         self.gw_interfaces = cast(
-            dict[str, BaseGWInterface], nn.ModuleDict(gw_interfaces)
+            dict[str, GWInterfaceBase], nn.ModuleDict(gw_interfaces)
         )
         self.workspace_dim = workspace_dim
 
@@ -234,7 +234,7 @@ def cycle(
         ...
 
 
-class GWInterface(BaseGWInterface):
+class GWInterface(GWInterfaceBase):
     def __init__(
         self,
         domain_module: DomainModule,
@@ -309,7 +309,7 @@ def cycle(
         }
 
 
-class VariationalGWInterface(BaseGWInterface):
+class VariationalGWInterface(GWInterfaceBase):
     def __init__(
         self,
         domain_module: DomainModule,