Update the fusion script and add the bayesian GW models

.vscode/launch.json

@@ -73,6 +73,18 @@
                 "CUDA_DEVICE_ORDER": "PCI_BUS_ID"
             }
         },
+        {
+            "type": "python",
+            "request": "launch",
+            "name": "Exp test common space fusion",
+            "program": "${workspaceFolder}/playground/exploration/test_common_space_fusion.py",
+            "justMyCode": false,
+            "env": {
+                "DEBUG": "1",
+                "CUDA_VISIBLE_DEVICES": "1",
+                "CUDA_DEVICE_ORDER": "PCI_BUS_ID"
+            }
+        },
         {
             "type": "python",
             "request": "launch",

config/default.yaml

@@ -13,7 +13,7 @@ training:
         max_lr: 5e-3
         weight_decay: 1e-5
 
-    max_steps: 15_000
+    max_steps: 200000
 
 logging:
     log_val_medias_every_n_epochs: 10
@@ -26,7 +26,7 @@ domain_modules:
         color_blind: false
 
 global_workspace:
-    has_uncertainty: false
+    bayesian_gw: false
 
     domain_proportions: 
         -   domains: ["v"]

config/exp_test_common_space_fusion.yaml

@@ -0,0 +1,2 @@
+exploration:
+    gw_checkpoint: "{default_root_dir}/shimmer-simple-shapes-t3p8nyq8/epoch=204.ckpt"

config/exp_test_t.yaml

@@ -1,2 +1,2 @@
 exploration:
-    gw_checkpoint: shimmer-simple-shapes-syhl3yoh/epoch=408.ckpt
+    gw_checkpoint: "{default_root_dir}/shimmer-simple-shapes-syhl3yoh/epoch=408.ckpt"

config/train_gw.yaml

@@ -4,7 +4,7 @@ training:
         max_lr: 5e-3
         weight_decay: 1e-6
 
-    max_steps: 100_000
+    max_steps: 50000
 
 logging:
     log_val_medias_every_n_epochs: 10

playground/exploration/contrastive_var.py

@@ -4,8 +4,8 @@
 
 import torch
 from shimmer import (
-    GlobalWorkspaceWithUncertainty,
-    GWLossesWithUncertainty,
+    GlobalWorkspaceBayesian,
+    GWLossesBayesian,
 )
 
 from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
@@ -85,7 +85,7 @@ def main():
 
     ckpt_path = config.default_root_dir / config.exploration.gw_checkpoint
     migrate_model(ckpt_path, PROJECT_DIR / "migrations" / "gw")
-    domain_module = GlobalWorkspaceWithUncertainty.load_from_checkpoint(
+    domain_module = GlobalWorkspaceBayesian.load_from_checkpoint(
         ckpt_path,
         domain_mods=domain_description,
         gw_encoders=gw_encoders,
@@ -122,7 +122,7 @@ def main():
             "attr": attr2.reshape(batch_size * n_rep, -1),
         }
     )
-    gw_states_std = gw_mod.log_uncertainties
+    gw_states_std = gw_mod.precisions
 
     actual_std_attr = (
         gw_states_means["attr"].reshape(batch_size, n_rep, -1).std(dim=1).mean(dim=0)
@@ -142,9 +142,7 @@ def main():
     print(f"Predicted std attr: {predicted_std_attr}")
     print(f"Predicted std v: {predicted_std_v}")
 
-    contrastive_fn = cast(
-        GWLossesWithUncertainty, domain_module.loss_mod
-    ).contrastive_fn
+    contrastive_fn = cast(GWLossesBayesian, domain_module.loss_mod).contrastive_fn
     assert contrastive_fn is not None
 
     norm1 = 1.0 + gw_states_std["attr"].exp() + gw_states_std["v_latents"].exp()

playground/exploration/norm_cont_comp.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from shimmer import GlobalWorkspaceWithUncertainty, contrastive_loss
+from shimmer import GlobalWorkspaceBayesian, contrastive_loss
 
 from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
 from simple_shapes_dataset.ckpt_migrations import migrate_model
@@ -80,7 +80,7 @@ def main():
 
     ckpt_path = config.default_root_dir / config.exploration.gw_checkpoint
     migrate_model(ckpt_path, PROJECT_DIR / "migrations" / "gw")
-    domain_module = GlobalWorkspaceWithUncertainty.load_from_checkpoint(
+    domain_module = GlobalWorkspaceBayesian.load_from_checkpoint(
         ckpt_path,
         domain_mods=domain_description,
         gw_encoders=gw_encoders,

playground/exploration/test_common_space_fusion.py

@@ -0,0 +1,107 @@
+import logging
+from collections.abc import Callable
+from typing import Any, cast
+
+from lightning.pytorch import Callback, Trainer, seed_everything
+from lightning.pytorch.callbacks import (
+    RichProgressBar,
+)
+from shimmer.modules.global_workspace import (
+    GlobalWorkspaceBayesian,
+    GWPredictionsBase,
+)
+from torch import set_float32_matmul_precision
+
+from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
+from simple_shapes_dataset.config import load_config
+from simple_shapes_dataset.dataset import SimpleShapesDataModule
+from simple_shapes_dataset.dataset.pre_process import (
+    color_blind_visual_domain,
+    nullify_attribute_rotation,
+)
+from simple_shapes_dataset.modules.domains import load_pretrained_domains
+
+
+def main():
+    config = load_config(
+        PROJECT_DIR / "config",
+        load_files=["exp_test_common_space_fusion.yaml"],
+        debug_mode=DEBUG_MODE,
+    )
+
+    if config.exploration is None:
+        raise ValueError("Config value 'exploration' must be set")
+
+    seed_everything(config.seed, workers=True)
+
+    domain_proportion = {
+        frozenset(item.domains): item.proportion
+        for item in config.global_workspace.domain_proportions
+    }
+
+    additional_transforms: dict[str, list[Callable[[Any], Any]]] = {}
+    if config.domain_modules.attribute.nullify_rotation:
+        logging.info("Nullifying rotation in the attr domain.")
+        additional_transforms["attr"] = [nullify_attribute_rotation]
+    if config.domain_modules.visual.color_blind:
+        logging.info("v domain will be color blind.")
+        additional_transforms["v"] = [color_blind_visual_domain]
+
+    data_module = SimpleShapesDataModule(
+        config.dataset.path,
+        domain_proportion,
+        batch_size=config.training.batch_size,
+        num_workers=config.training.num_workers,
+        seed=config.seed,
+        ood_seed=config.ood_seed,
+        domain_args=config.global_workspace.domain_args,
+        additional_transforms=additional_transforms,
+    )
+
+    domain_modules, gw_encoders, gw_decoders = load_pretrained_domains(
+        config.default_root_dir,
+        config.global_workspace.domains,
+        config.global_workspace.latent_dim,
+        config.global_workspace.encoders.hidden_dim,
+        config.global_workspace.encoders.n_layers,
+        config.global_workspace.decoders.hidden_dim,
+        config.global_workspace.decoders.n_layers,
+        is_linear=config.global_workspace.linear_domains,
+        bias=config.global_workspace.linear_domains_use_bias,
+    )
+
+    module = GlobalWorkspaceBayesian.load_from_checkpoint(
+        config.exploration.gw_checkpoint,
+        domain_mods=domain_modules,
+        gw_encoders=gw_encoders,
+        gw_decoders=gw_decoders,
+    )
+
+    callbacks: list[Callback] = [RichProgressBar()]
+    set_float32_matmul_precision(config.training.float32_matmul_precision)
+
+    trainer = Trainer(
+        fast_dev_run=config.training.fast_dev_run,
+        max_steps=config.training.max_steps,
+        enable_progress_bar=config.training.enable_progress_bar,
+        default_root_dir=config.default_root_dir,
+        callbacks=callbacks,
+        precision=config.training.precision,
+        accelerator=config.training.accelerator,
+        devices=config.training.devices,
+    )
+
+    predictions = cast(
+        list[GWPredictionsBase],
+        trainer.predict(module, data_module, return_predictions=True),
+    )
+
+    for pred in predictions:
+        for k in range(pred["states"]["attr"].size(0)):
+            for domain in pred["states"]:
+                print(domain, pred["states"][domain][k])
+            print("--")
+
+
+if __name__ == "__main__":
+    main()

playground/exploration/var_norm_cont.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 from lightning.pytorch import Trainer
-from shimmer.modules.global_workspace import GlobalWorkspaceWithUncertainty
+from shimmer.modules.global_workspace import GlobalWorkspaceBayesian
 
 from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
 from simple_shapes_dataset.ckpt_migrations import (
@@ -63,7 +63,7 @@ def main():
 
     ckpt_path = config.exploration.gw_checkpoint
     migrate_model(ckpt_path, PROJECT_DIR / "migrations" / "gw")
-    gw = GlobalWorkspaceWithUncertainty.load_from_checkpoint(
+    gw = GlobalWorkspaceBayesian.load_from_checkpoint(
         ckpt_path,
         domain_mods=domain_description,
         gw_encoders=gw_encoders,

playground/exploration/var_norm_cont_loss.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from shimmer.modules.global_workspace import GlobalWorkspaceWithUncertainty
+from shimmer.modules.global_workspace import GlobalWorkspaceBayesian
 
 from simple_shapes_dataset import DEBUG_MODE, PROJECT_DIR
 from simple_shapes_dataset.ckpt_migrations import (
@@ -82,7 +82,7 @@ def main():
 
     ckpt_path = config.default_root_dir / config.exploration.gw_checkpoint
     migrate_model(ckpt_path, PROJECT_DIR / "migrations" / "gw")
-    domain_module = GlobalWorkspaceWithUncertainty.load_from_checkpoint(
+    domain_module = GlobalWorkspaceBayesian.load_from_checkpoint(
         ckpt_path,
         domain_mods=domain_description,
         gw_encoders=gw_encoders,
@@ -103,7 +103,7 @@ def main():
     v_test[:, :12] = v_paired[None, :12]
     attr_test[:, :12] = attr_paired[None, :12]
     gw_states_means = gw_mod.encode({"v_latents": v_test, "attr": attr_test})
-    gw_states_std = gw_mod.log_uncertainties
+    gw_states_std = gw_mod.precisions
     v_gw_var = (0.5 * gw_states_std["v_latents"]).exp()  # noqa: F841
     attr_gw_var = (0.5 * gw_states_std["attr"]).exp()  # noqa: F841
     print(gw_states_means)

playground/train_gw.py

@@ -10,11 +10,17 @@
     RichProgressBar,
 )
 from lightning.pytorch.loggers.wandb import WandbLogger
-from shimmer import ContrastiveLossType, GlobalWorkspaceBase, LossCoefs, SaveMigrations
+from shimmer import (
+    BroadcastLossCoefs,
+    ContrastiveLossType,
+    GlobalWorkspaceBase,
+    LossCoefs,
+    SaveMigrations,
+)
 from shimmer.modules.global_workspace import (
     GlobalWorkspace,
-    GlobalWorkspaceFusion,
-    GlobalWorkspaceWithUncertainty,
+    GlobalWorkspace2Domains,
+    GlobalWorkspaceBayesian,
     SchedulerArgs,
 )
 from torch import set_float32_matmul_precision
@@ -76,13 +82,6 @@ def main():
         bias=config.global_workspace.linear_domains_use_bias,
     )
 
-    loss_coefs: LossCoefs = {
-        "demi_cycles": config.global_workspace.loss_coefficients.demi_cycles,
-        "cycles": config.global_workspace.loss_coefficients.cycles,
-        "translations": config.global_workspace.loss_coefficients.translations,
-        "contrastives": config.global_workspace.loss_coefficients.contrastives,
-    }
-
     contrastive_fn: ContrastiveLossType | None = None
     if config.global_workspace.vsepp_contrastive_loss:
         contrastive_fn = VSEPPContrastiveLoss(
@@ -93,13 +92,23 @@ def main():
         )
 
     module: GlobalWorkspaceBase
-    if config.global_workspace.has_uncertainty:
-        module = GlobalWorkspaceWithUncertainty(
+    gw_type: str
+    if config.global_workspace.bayesian_gw:
+        gw_type = "gw_bayesian"
+        loss_coefs_bayesian: BroadcastLossCoefs = {
+            "contrastives": config.global_workspace.loss_coefficients.contrastives,
+            "fused": config.global_workspace.loss_coefficients.fused,
+            "translations": config.global_workspace.loss_coefficients.translations,
+            "demi_cycles": config.global_workspace.loss_coefficients.demi_cycles,
+            "cycles": config.global_workspace.loss_coefficients.cycles,
+        }
+        module = GlobalWorkspaceBayesian(
             domain_modules,
             gw_encoders,
             gw_decoders,
             config.global_workspace.latent_dim,
-            loss_coefs,
+            loss_coefs_bayesian,
+            config.global_workspace.selection_temperature,
             config.training.optim.lr,
             config.training.optim.weight_decay,
             scheduler_args=SchedulerArgs(
@@ -110,11 +119,21 @@ def main():
             contrastive_loss=contrastive_fn,
         )
     elif config.global_workspace.use_fusion_model:
-        module = GlobalWorkspaceFusion(
+        gw_type = "gw_fusion"
+        loss_coefs_fusion: BroadcastLossCoefs = {
+            "contrastives": config.global_workspace.loss_coefficients.contrastives,
+            "fused": config.global_workspace.loss_coefficients.fused,
+            "translations": config.global_workspace.loss_coefficients.translations,
+            "demi_cycles": config.global_workspace.loss_coefficients.demi_cycles,
+            "cycles": config.global_workspace.loss_coefficients.cycles,
+        }
+        module = GlobalWorkspace(
             domain_modules,
             gw_encoders,
             gw_decoders,
             config.global_workspace.latent_dim,
+            loss_coefs_fusion,
+            config.global_workspace.selection_temperature,
             config.training.optim.lr,
             config.training.optim.weight_decay,
             scheduler_args=SchedulerArgs(
@@ -125,7 +144,15 @@ def main():
             contrastive_loss=contrastive_fn,
         )
     else:
-        module = GlobalWorkspace(
+        gw_type = "gw"
+        loss_coefs: LossCoefs = {
+            "demi_cycles": config.global_workspace.loss_coefficients.demi_cycles,
+            "cycles": config.global_workspace.loss_coefficients.cycles,
+            "translations": config.global_workspace.loss_coefficients.translations,
+            "contrastives": config.global_workspace.loss_coefficients.contrastives,
+        }
+
+        module = GlobalWorkspace2Domains(
             domain_modules,
             gw_encoders,
             gw_decoders,
@@ -216,7 +243,6 @@ def main():
 
     wandb_logger = None
     if config.wandb.enabled:
-        gw_type = "gw_uncertainty" if config.global_workspace.has_uncertainty else "gw"
         run_name = f"{gw_type}_z={config.global_workspace.latent_dim}"
         wandb_logger = WandbLogger(
             save_dir=config.wandb.save_dir,