simple-shapes-dataset-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "72c6c049-fa9d-4add-8156-7ec359734532",
   "metadata": {},
   "source": [
    "# How to train your Global Workspace\n",
    "Benjamin Devillers\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ruflab/shimmer-tutorials/blob/main/simple-shapes-dataset-training.ipynb)\n",
    "\n",
    "\n",
    "In this notebook, we will see how to use `shimmer` to build and train from scratch a Global Workspace on the Simple Shapes Dataset. We train a model than can translate visual images of shapes from the [simple-shapes-datset](https://github.com/ruflab/simple-shapes-dataset) to their proto-language (attributes)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e7893c3-98e4-47ba-96bb-13f2652e45c6",
   "metadata": {},
   "source": [
    "For this tutorial, we will need to install the [shimmer-ssd](https://github.com/ruflab/shimmer-ssd) package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e2a857-f0a5-4b98-aa95-2e12bf15fb87",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"git+https://github.com/ruflab/shimmer-ssd.git\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7fef0c8c-3381-4b8c-896f-1ae2b443db61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting tensorboardX\n",
      "  Downloading tensorboardX-2.6.2.2-py2.py3-none-any.whl.metadata (5.8 kB)\n",
      "Requirement already satisfied: numpy in /home/bdevillers/.cache/pypoetry/virtualenvs/shimmer-tutorials-RRvxJ_Ue-py3.11/lib/python3.11/site-packages (from tensorboardX) (1.26.4)\n",
      "Requirement already satisfied: packaging in /home/bdevillers/.cache/pypoetry/virtualenvs/shimmer-tutorials-RRvxJ_Ue-py3.11/lib/python3.11/site-packages (from tensorboardX) (24.2)\n",
      "Requirement already satisfied: protobuf>=3.20 in /home/bdevillers/.cache/pypoetry/virtualenvs/shimmer-tutorials-RRvxJ_Ue-py3.11/lib/python3.11/site-packages (from tensorboardX) (5.29.3)\n",
      "Downloading tensorboardX-2.6.2.2-py2.py3-none-any.whl (101 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m101.7/101.7 kB\u001b[0m \u001b[31m2.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m \u001b[36m0:00:01\u001b[0m\n",
      "Installing collected packages: tensorboardX\n",
      "Successfully installed tensorboardX-2.6.2.2\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install tensorboard"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17deebf1-86e1-46db-bc91-820158d36bcc",
   "metadata": {},
   "source": [
    "This package depends on [simple-shapes-dataset](https://github.com/ruflab/simple-shapes-dataset) and provides all of its commands. You can then use all of its commands.\n",
    "\n",
    "For instance, we can download the dataset directly with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80bf4524-db6c-439c-a491-0e5883b0861c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!shapesd download"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29130f4-ac61-460c-bc7f-cd0f7437fd5e",
   "metadata": {},
   "source": [
    "Note that `shapesd download` automatically migrates the dataset so that it is correctly formatted. If you downloaded the dataset manually, use `shapesd migrate -p PATH_TO_DATASET` to migrate manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "75a2017e-7de7-4568-87dd-23bf4f93a0cc",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from collections.abc import Mapping, Sequence\n",
    "from pathlib import Path\n",
    "from typing import Any, cast\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from lightning.pytorch import Callback, Trainer, seed_everything\n",
    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
    "from lightning.pytorch.loggers import TensorBoardLogger\n",
    "from shimmer import DomainModule, LossOutput\n",
    "from shimmer.modules.domain import DomainModule\n",
    "from shimmer.modules.global_workspace import GlobalWorkspace2Domains, SchedulerArgs\n",
    "from shimmer.modules.vae import (\n",
    "    VAE,\n",
    "    VAEDecoder,\n",
    "    VAEEncoder,\n",
    "    gaussian_nll,\n",
    "    kl_divergence_loss,\n",
    ")\n",
    "from shimmer_ssd import DEBUG_MODE, LOGGER, PROJECT_DIR\n",
    "from shimmer_ssd.config import DomainModuleVariant, LoadedDomainConfig, load_config\n",
    "from shimmer_ssd.dataset.pre_process import TokenizeCaptions\n",
    "from shimmer_ssd.logging import (\n",
    "    LogAttributesCallback,\n",
    "    LogGWImagesCallback,\n",
    "    LogVisualCallback,\n",
    "    batch_to_device,\n",
    ")\n",
    "from shimmer_ssd.modules.domains import load_pretrained_domains\n",
    "from shimmer_ssd.modules.domains.visual import VisualLatentDomainModule\n",
    "from shimmer_ssd.modules.vae import RAEDecoder, RAEEncoder\n",
    "from tokenizers.implementations.byte_level_bpe import ByteLevelBPETokenizer\n",
    "from torch import nn\n",
    "from torch.nn.functional import mse_loss\n",
    "from torch.optim.lr_scheduler import OneCycleLR\n",
    "from torch.optim.optimizer import Optimizer\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "from simple_shapes_dataset import SimpleShapesDataModule, get_default_domains\n",
    "\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8d9c35d-7b83-4ab3-9da2-9bc220a5ad87",
   "metadata": {},
   "source": [
    "## Config\n",
    "\n",
    "Let's first generate the config folder for the rest of the scripts.\n",
    "This will create a `config` folder with different yaml files used by the different scripts and in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b9272822-413b-4580-bf0e-93a117112e67",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config folder already exists. Skipping.\n"
     ]
    }
   ],
   "source": [
    "!ssd config create"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01bd8041-8e59-429e-9c6c-e8b656ecbbb7",
   "metadata": {},
   "source": [
    "This will create a `config` folder. This contains many file, but in this tutorial, only `main.yaml` will interest us.\n",
    "\n",
    "You can start by taking a look at the default values which should be mostly set correctly for this tutorial. But you can try and make some changes to see the outcome.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "Anytime you make a change to the config, don't forget to reload it with the following cell!\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e7246f6e-a4ac-4e01-9778-d75021c698e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We don't use cli in the notebook, but consider using it in normal scripts.\n",
    "config = load_config(\"./config\", use_cli=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15df99ca-896c-4443-8ed4-da143578aa50",
   "metadata": {},
   "source": [
    "## Data format\n",
    "\n",
    "The dataloader provides the data in a specific format:\n",
    "\n",
    "```python\n",
    "domain_group = {\n",
    "    \"domain\": domain_data\n",
    "}\n",
    "batch = {\n",
    "    frozenset([\"domain\"]): domain_group\n",
    "}\n",
    "```\n",
    "* The **batch** is a dict that has frozensets of domains as keys, and a domain group as values.\n",
    "* The **domain group** is a dict that has domains (string) as keys, and the domain data as values. The data samples of every domain in a domain group is matched. This\n",
    "means that for a domain group that has 2 domains d1 and d2: `domain_group[\"d1\"][k]` is paired with `domain_group[\"d2\"][k]` for all `k`.\n",
    "\n",
    "This allows a batch to have several groups (of different domains) of paired data. For example, a batch with unpaired visual (domain \"v\"), unpaired attribute (domain \"attr\"), and paired visual and attribute will look like:\n",
    "```python\n",
    "batch = {\n",
    "    frozenset([\"v\"]): {\"v\": unpaired_visual_data},\n",
    "    frozenset([\"attr\"]): {\"attr\": unpaired_attribute_data},\n",
    "    frozenset([\"attr\", \"v\"]): {\"attr\": paired_attr_data, \"v\": paired_visual_data},\n",
    "}\n",
    "```\n",
    "\n",
    "This is useful to train the global workspace later. But this is also the format used to train the unimodal domains.\n",
    "\n",
    "Note that because all the data is paired in validation and test steps, the dataloader only returns one domain group with all paired domain:\n",
    "```python\n",
    "val_batch = {\"attr\": paired_attr_data, \"v\": paired_v_data}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32069c2-3867-47f5-a72e-f6f077895e40",
   "metadata": {},
   "source": [
    "## Unimodal Domains\n",
    "\n",
    "Instead of using the domains already defined in [`shimmer-ssd`](https://ruflab.github.io/shimmer-ssd/latest/), we will redefine them from scratch to learn how to make them.\n",
    "\n",
    "We will train 2 domain modules:\n",
    "* vision: a beta-VAE that will encode the shape images into a small latent representation vector;\n",
    "* attribute: a beta-VAE that will encode the object category and other attributes into a latent representation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee05e0db-512b-4888-8496-b73169eb088e",
   "metadata": {},
   "source": [
    "### Vision domain\n",
    "\n",
    "Let's start by defining the data module to train the vision module.\n",
    "\n",
    "We will load the config with an extra argument that will load the `train_v.yaml` file. This file contains specific configuration related to training the visual side like learning rate, weight decay, and the number of steps to train for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "36ed5cb2-8e63-4d1e-a3c1-7987dedc0535",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config(\"./config\", use_cli=False, load_files=[\"train_v.yaml\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d38e8ada-6293-4cd7-9a7a-c3fbcde6989c",
   "metadata": {},
   "source": [
    "Note that with `load_files=[\"train_v.yaml\"]`, the `main.yaml` file is still loaded and you don't have to repeat it. If you repeat some config values, the settings set in `train_v.yaml` will override the values in `main.yaml`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c3af7304-b507-4981-95ab-128daa126b53",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 0\n"
     ]
    }
   ],
   "source": [
    "seed_everything(config.seed, workers=True)\n",
    "\n",
    "data_module = SimpleShapesDataModule(\n",
    "    config.dataset.path,\n",
    "    get_default_domains([\"v\"]),\n",
    "    {frozenset([\"v\"]): 1.0},\n",
    "    batch_size=config.training.batch_size,  # set in `config/train_v.yaml`\n",
    "    num_workers=config.training.num_workers,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ee9b55a-41fb-4e4b-9a6a-ebdb6dd01af2",
   "metadata": {},
   "source": [
    "All domain modules must inherit from the class `DomainModule` from `shimmer`. This class requires the creation of several methods.\n",
    "`DomainModule`s inherit from a `LightningModule`, so you can also define the relevant method to train it.\n",
    "\n",
    "In particular, you will need to overide the following methods:\n",
    "* `def encode(self, x: Any) -> torch.Tensor`: encodes the raw domain data into the unimodal latent representation\n",
    "* `def decode(self, z: torch.Tensor) -> Any`: decodes the unimodal latent representation into the raw input data\n",
    "* `def compute_loss(self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any) -> LossOutput | None`: how to compute the loss used for cycle constency or translation\n",
    "of this domain. `pred` is the predicted unimodal latent vector, `target` is the target latent vector and `raw_target` is the original input before being encoded.\n",
    "\n",
    "You can see more details in [the docs](https://ruflab.github.io/shimmer/latest/shimmer/modules/domain.html#DomainModule).\n",
    "\n",
    "In addition, we will also define `configure_optimizers`, `training_step`, and `validation_step`, to define the optimizer, and what happens during the training and validation steps. See the [lightning docs](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b21d1398-c487-420e-8b31-3844e12a66a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class VisualDomainModule(DomainModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_channels: int,\n",
    "        latent_dim: int,\n",
    "        ae_dim: int,\n",
    "        beta: float = 1,\n",
    "        optim_lr: float = 1e-3,\n",
    "        optim_weight_decay: float = 0,\n",
    "        scheduler_args: Mapping[str, Any] | None = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Visual domain module. This defines shimmer's `DomainModule` for the vision\n",
    "        side with a VAE.\n",
    "\n",
    "        Args:\n",
    "            num_channels (`int`): number of input channels (for RGB image, use 3)\n",
    "            latent_dim (`int`): latent dimension of the vision domain\n",
    "            ae_dim (`int`): internal auto-encoder dimension of the VAE\n",
    "            beta (`float`): beta value if beta-VAE. (Defaults to 1.0)\n",
    "            optim_lr (`float`): training learning rate\n",
    "            optim_weight_decay (`float`): training weight decay\n",
    "            scheduler_args (`Mapping[str, Any] | None`): Args for the scheduler.\n",
    "        \"\"\"\n",
    "\n",
    "        # The parent class requires to know the latent dimension of the module.\n",
    "        super().__init__(latent_dim)\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        vae_encoder = RAEEncoder(num_channels, ae_dim, latent_dim, use_batchnorm=True)\n",
    "        vae_decoder = RAEDecoder(num_channels, latent_dim, ae_dim)\n",
    "\n",
    "        # This is a useful helper to make a VAE. You can learn more on this here if interested:\n",
    "        # https://ruflab.github.io/shimmer-ssd/latest/shimmer_ssd/modules/vae.html\n",
    "        self.vae = VAE(vae_encoder, vae_decoder, beta)\n",
    "        self.optim_lr = optim_lr\n",
    "        self.optim_weight_decay = optim_weight_decay\n",
    "        self.scheduler_args: dict[str, Any] = {\n",
    "            \"max_lr\": optim_lr,\n",
    "            \"total_steps\": 1,\n",
    "        }\n",
    "        self.scheduler_args.update(scheduler_args or {})\n",
    "\n",
    "    def encode(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Encode from the image to the latent representation.\n",
    "        Here we can just use the encode function from the VAE.\n",
    "        \"\"\"\n",
    "        return self.vae.encode(x)\n",
    "\n",
    "    def decode(self, z: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Decode the unimodal latent into the original domain.\n",
    "        We can use the decode function from the VAE.\n",
    "        \"\"\"\n",
    "        return self.vae.decode(z)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        return self.decode(self.encode(x))\n",
    "\n",
    "    def compute_loss(self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any) -> LossOutput:\n",
    "        \"\"\"\n",
    "        Compute MSE loss in the latent domain. This will be usefull when training the global workspace.\n",
    "        This defines the loss to use for the cycle consistency losses and the translation losses.\n",
    "        \"\"\"\n",
    "        loss = mse_loss(pred, target, reduction=\"mean\")\n",
    "        return LossOutput(loss)\n",
    "\n",
    "    # Pytorch Lightning related methods\n",
    "\n",
    "    def training_step(self, batch: Mapping[frozenset[str], Mapping[str, torch.Tensor]], batch_idx: int) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            batch (`Mapping[frozenset[str], Mapping[str, torch.Tensor]]`): batch of domain groups. When training the visual domain,\n",
    "                this will only contain one group with unpaired visual data\n",
    "            batch_idx (`int`): batch index\n",
    "        Returns:\n",
    "            `torch.Tensor`: the total loss of the step\n",
    "        \"\"\"\n",
    "        x = batch[frozenset([\"v\"])][\"v\"]  # extracts the visual information\n",
    "        return self.generic_step(x, \"train\")  # use a generic step method that handles both train and validation.\n",
    "\n",
    "    def validation_step(self, batch: Mapping[str, torch.Tensor], batch_idx: int) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Validation step\n",
    "\n",
    "        Args:\n",
    "            batch (`Mapping[str, torch.Tensor]`): group with only paired data. The validation step does not receive several group.\n",
    "            batch_idx (`int`): batch index\n",
    "\n",
    "            Returns:\n",
    "                `torch.Tensor`: the total loss\n",
    "        \"\"\"\n",
    "        x = batch[\"v\"]  # extract the visual data of the group\n",
    "        return self.generic_step(x, \"val\")  # use `generic_step` method to compute the val loss.\n",
    "\n",
    "    def generic_step(self, x: torch.Tensor, mode: str = \"train\") -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Computes the loss given image data\n",
    "\n",
    "        Args:\n",
    "            x (`torch.Tensor`): tensors of images\n",
    "            mode (`str`): mode of step (\"train\", \"val\") used for logging\n",
    "        Returns:\n",
    "            `torch.Tensor`: The computed loss\n",
    "        \"\"\"\n",
    "        # Get the latent mean, log-variance, and reconstructed images\n",
    "        # from the VAE\n",
    "        (mean, logvar), reconstruction = self.vae(x)\n",
    "\n",
    "        reconstruction_loss = gaussian_nll(reconstruction, torch.tensor(0), x).sum()\n",
    "        kl_loss = kl_divergence_loss(mean, logvar)\n",
    "        total_loss = reconstruction_loss + self.vae.beta * kl_loss\n",
    "\n",
    "        # Log losses\n",
    "        self.log(f\"{mode}/reconstruction_loss\", reconstruction_loss)\n",
    "        self.log(f\"{mode}/kl_loss\", kl_loss)\n",
    "        self.log(f\"{mode}/loss\", total_loss)\n",
    "\n",
    "        return total_loss\n",
    "\n",
    "    def configure_optimizers(self) -> dict[str, Any]:\n",
    "        \"\"\"\n",
    "        Defines optimizer and learning rate scheduler\n",
    "        \"\"\"\n",
    "\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            self.parameters(),\n",
    "            lr=self.optim_lr,\n",
    "            weight_decay=self.optim_weight_decay,\n",
    "        )\n",
    "        lr_scheduler = OneCycleLR(optimizer, **self.scheduler_args)\n",
    "\n",
    "        return {\n",
    "            \"optimizer\": optimizer,\n",
    "            \"lr_scheduler\": {\n",
    "                \"scheduler\": lr_scheduler,\n",
    "                \"interval\": \"step\",\n",
    "            },\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71b37122-2562-4daf-b9bc-7f98a5022862",
   "metadata": {},
   "source": [
    "Now let's instanciate the domain module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bf4977c2-7dca-4a2f-9dea-525279b503d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_domain_module = VisualDomainModule(\n",
    "    num_channels=3,\n",
    "    ae_dim=config.domain_modules.visual.ae_dim,\n",
    "    latent_dim=config.domain_modules.visual.latent_dim,\n",
    "    beta=config.domain_modules.visual.beta,\n",
    "    optim_lr=config.training.optim.lr,\n",
    "    optim_weight_decay=config.training.optim.weight_decay,\n",
    "    scheduler_args={\n",
    "        \"max_lr\": config.training.optim.max_lr,\n",
    "        \"total_steps\": config.training.max_steps,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19ebfe33-58ab-4e9b-ac77-ef9018c7a301",
   "metadata": {},
   "source": [
    "We will use tensorboard to log the losses and reconstructed images. We can use `LogVisualCallback` from `shimmer_ssd.logging` to\n",
    "log reconstructed images of some image samples.\n",
    "\n",
    "You can update the `train_v.yaml` config file to change how often images will be updated on tensorboard. Here we will put 1 so that they are updated every epoch:\n",
    "```yaml\n",
    "logging:\n",
    "    log_val_medias_every_n_epochs: 1\n",
    "    log_train_medias_every_n_epochs: 1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "008f09cf-86bf-4a14-8a64-a17c00723ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = TensorBoardLogger(\"logs\", name=\"vision_model\")\n",
    "\n",
    "# Get some image samples to log in tensorboard.\n",
    "val_samples = data_module.get_samples(\"val\", 32)[frozenset([\"v\"])][\"v\"]\n",
    "train_samples = data_module.get_samples(\"train\", 32)[frozenset([\"v\"])][\"v\"]\n",
    "\n",
    "# Create vision vision where we will save checkpoints\n",
    "(config.default_root_dir / \"vision\").mkdir(exist_ok=True)\n",
    "\n",
    "callbacks: list[Callback] = [\n",
    "    # Will log the validation ground-truth and reconstructions during training\n",
    "    LogVisualCallback(\n",
    "        val_samples,\n",
    "        log_key=\"val_images\",\n",
    "        mode=\"val\",\n",
    "        every_n_epochs=config.logging.log_val_medias_every_n_epochs,\n",
    "        ncols=8,\n",
    "    ),\n",
    "    # Will log the training ground-truth and reconstructions during training\n",
    "    LogVisualCallback(\n",
    "        train_samples,\n",
    "        log_key=\"train_images\",\n",
    "        mode=\"train\",\n",
    "        every_n_epochs=config.logging.log_train_medias_every_n_epochs,\n",
    "        ncols=8,\n",
    "    ),\n",
    "    # Save the checkpoints\n",
    "    ModelCheckpoint(\n",
    "        dirpath=config.default_root_dir / \"vision\" / f\"version_{logger.version}\",\n",
    "        filename=\"{epoch}\",\n",
    "        monitor=\"val/loss\",\n",
    "        mode=\"min\",\n",
    "        save_last=\"link\",\n",
    "        save_top_k=1,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d3dbe1a-fbaa-4f7f-b2e0-2bb26d4eee71",
   "metadata": {},
   "source": [
    "For the final model, let's save where the model is saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0f3b5b99-5884-41c0-96cd-15a212331d4a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoints/vision/version_2\n"
     ]
    }
   ],
   "source": [
    "visual_checkpoint = config.default_root_dir / \"vision\" / f\"version_{logger.version}\"\n",
    "print(visual_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aee15178-da03-43cf-b90a-ebd7265a1df6",
   "metadata": {},
   "source": [
    "Load tensorboard. You can select the version associated to the previous path. It will appear after the training is started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "27fa7d23-e66d-4d49-98f1-5bff9b3085fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "153ef650-82ac-44bf-a964-442ea7254a7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-629f6fbed82c07cd\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-629f6fbed82c07cd\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir \"./logs/vision_model\" --reload_interval 30 --reload_task 'auto'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7493459-eba1-421a-97f4-0a718bb1bce3",
   "metadata": {},
   "source": [
    "Let'st start the training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7ebadf89-0c0c-46e7-9dc8-52d75db37537",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "You are using a CUDA device ('NVIDIA A100 80GB PCIe') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "\n",
      "  | Name | Type | Params | Mode \n",
      "--------------------------------------\n",
      "0 | vae  | VAE  | 1.5 M  | train\n",
      "--------------------------------------\n",
      "1.5 M     Trainable params\n",
      "0         Non-trainable params\n",
      "1.5 M     Total params\n",
      "5.990     Total estimated model params size (MB)\n",
      "32        Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68fadfa5b5a5426a84c2b3c1bfa39a18",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |                                                                                            …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb8bf278e1ab4ca5a98cc9aeea400871",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |                                                                                                   …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85fe9ba9075e4cc89305b71cf8fa8bcb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "32bfd005db2648b7babe47509827f940",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d81198448b7c48d2bde1c45055c29810",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ff45aad7e2844ba81c121e3c8ac2e63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "798acefbc66841699d97070e988b27d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b9f8ba444d647aa92d085e8904949aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17d864fc8391479bbcc4b83429a56f48",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4171765bb4ed48c2a852a2a88bce52d4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_steps=2000` reached.\n",
      "Restoring states from the checkpoint path at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/vision/version_2/epoch=7.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "Loaded model weights from the checkpoint at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/vision/version_2/epoch=7.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "16830b2db614452eaf3f52d8d110f4a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">      Validate metric      </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">        val/kl_loss        </span>│<span style=\"color: #800080; text-decoration-color: #800080\">       37061.0546875       </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">         val/loss          </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         5761903.5         </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">  val/reconstruction_loss  </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         5758197.0         </span>│\n",
       "└───────────────────────────┴───────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m     Validate metric     \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m       val/kl_loss       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m      37061.0546875      \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m        val/loss         \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        5761903.5        \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m val/reconstruction_loss \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        5758197.0        \u001b[0m\u001b[35m \u001b[0m│\n",
       "└───────────────────────────┴───────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'val/reconstruction_loss': 5758197.0,\n",
       "  'val/kl_loss': 37061.0546875,\n",
       "  'val/loss': 5761903.5}]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    logger=logger,\n",
    "    fast_dev_run=config.training.fast_dev_run,\n",
    "    max_steps=config.training.max_steps,\n",
    "    enable_progress_bar=config.training.enable_progress_bar,\n",
    "    default_root_dir=config.default_root_dir,\n",
    "    callbacks=callbacks,\n",
    "    precision=config.training.precision,\n",
    "    accelerator=config.training.accelerator,\n",
    "    devices=config.training.devices,\n",
    ")\n",
    "\n",
    "trainer.fit(v_domain_module, data_module)\n",
    "trainer.validate(v_domain_module, data_module, \"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "287fc371-c0a0-4bd5-b3e4-5b81b2335c47",
   "metadata": {},
   "source": [
    "For faster training of the global workspace, we can extract the visual latent feature with the following cell. Don't forget to update the path to the actual checkpoint, which is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0032a359-eca8-4c8b-a1c0-d04b299ad9a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoints/vision/version_2\n"
     ]
    }
   ],
   "source": [
    "print(visual_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "323b2ece-eac5-45ff-842e-ae49e3070569",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Latent file already exists. Overriding.\n",
      "Saving train.\n",
      "100%|█████████████████████████████████████████| 244/244 [00:11<00:00, 22.02it/s]\n",
      "Saving in simple_shapes_dataset/saved_latents/train/domain_v_tuto.npy.\n",
      "Saving val.\n",
      "100%|███████████████████████████████████████████| 25/25 [00:01<00:00, 14.02it/s]\n",
      "Saving in simple_shapes_dataset/saved_latents/val/domain_v_tuto.npy.\n",
      "Saving test.\n",
      "Traceback (most recent call last):                       | 0/25 [00:00<?, ?it/s]\n",
      "  File \"/home/bdevillers/.pyenv/versions/3.11.4/lib/python3.11/multiprocessing/util.py\", line 345, in _run_finalizers\n",
      "    finalizer()\n",
      "  File \"/home/bdevillers/.pyenv/versions/3.11.4/lib/python3.11/multiprocessing/util.py\", line 266, in __call__\n",
      "    res = self._callback(*self._args, **self._kwargs)\n",
      "          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/home/bdevillers/.pyenv/versions/3.11.4/lib/python3.11/multiprocessing/util.py\", line 153, in _remove_temp_dir\n",
      "    rmtree(tempdir)\n",
      "  File \"/home/bdevillers/.pyenv/versions/3.11.4/lib/python3.11/shutil.py\", line 738, in rmtree\n",
      "    onerror(os.rmdir, path, sys.exc_info())\n",
      "  File \"/home/bdevillers/.pyenv/versions/3.11.4/lib/python3.11/shutil.py\", line 736, in rmtree\n",
      "    os.rmdir(path, dir_fd=dir_fd)\n",
      "OSError: [Errno 39] Directory not empty: '/tmp/pymp-a637noda'\n",
      "100%|███████████████████████████████████████████| 25/25 [00:01<00:00, 13.95it/s]\n",
      "Saving in simple_shapes_dataset/saved_latents/test/domain_v_tuto.npy.\n"
     ]
    }
   ],
   "source": [
    "!ssd extract v \"checkpoints/vision/version_2/last.ckpt\" -p \"simple_shapes_dataset\" --latent_name \"domain_v_tuto.npy\" --force"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a70c380-8033-4728-a6d3-3c32fe8ad0f7",
   "metadata": {},
   "source": [
    "We now need to update the location of the presaved latent vectors in the config file. You can see that `domain_data_args` is defined as follows:\n",
    "\n",
    "```yaml\n",
    "domain_data_args:\n",
    "    v_latents:\n",
    "        presaved_path: domain_v_tuto.npy  # as defined in the `--latent_name`\n",
    "```\n",
    "let's change the selected domains:\n",
    "\n",
    "```yaml\n",
    "domains:\n",
    "    - checkpoint_path: \"./checkpoints/visual/version_0/last.ckpt\"  # update to the actual version\n",
    "      domain_type: v_latents\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6c5ed2c-9fcd-42fa-a6a7-bd8d6c7a8ef8",
   "metadata": {},
   "source": [
    "### Attribute domain\n",
    "\n",
    "This will be very similar to the previous section. Here, we will focus on learning the attribute model.\n",
    "\n",
    "We will load the config with an extra argument that will load the `train_attr.yaml` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e85d14bf-5406-4e2b-bd82-c9936abb039c",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config(\"./config\", use_cli=False, load_files=[\"train_attr.yaml\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1dbd937-2030-4345-a648-54456cdb99e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Seed set to 0\n"
     ]
    }
   ],
   "source": [
    "seed_everything(config.seed, workers=True)\n",
    "\n",
    "data_module = SimpleShapesDataModule(\n",
    "    config.dataset.path,\n",
    "    get_default_domains([\"attr\"]),\n",
    "    {frozenset([\"attr\"]): 1.0},\n",
    "    batch_size=config.training.batch_size,  # set in `config/train_attr.yaml`\n",
    "    num_workers=config.training.num_workers,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6177d88f-bce3-4393-82a2-02fab9802c91",
   "metadata": {},
   "source": [
    "Similarly to what has been done for the visual domain, we will create a `DomainModule`. This will also use a VAE to encode and decode the attribute vectors.\n",
    "\n",
    "First, we will define the VAE encoders and decoders. They will inherit from `shimme.modules.vae.VAEEncoder` and `shimme.modules.vae.VAEEncoder`.\n",
    "\n",
    "For the encoder:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c0db8a68-28a1-4080-bdad-8b41ff7cb09f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(VAEEncoder):\n",
    "    def __init__(\n",
    "        self,\n",
    "        hidden_dim: int,\n",
    "        out_dim: int,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.out_dim = out_dim\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            # Input dim: 3 one-hot encoded shape category + 2 locations + 2 rotations (cos, sin space) + 1 size + 3 color (RGB)\n",
    "            nn.Linear(11, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, out_dim),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        self.q_mean = nn.Linear(self.out_dim, self.out_dim)\n",
    "        self.q_logvar = nn.Linear(self.out_dim, self.out_dim)\n",
    "\n",
    "    def forward(self, x: Sequence[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:\n",
    "        out = torch.cat(list(x), dim=-1)\n",
    "        out = self.encoder(out)\n",
    "        return self.q_mean(out), self.q_logvar(out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32d0d54a-c860-4bbe-a874-bc284bd19e19",
   "metadata": {},
   "source": [
    "And for the decoder, we will decode back to two vectors: a class vector, and an attribute one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b298e147-cc5a-44b2-97f5-be1e2d43b40c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(VAEDecoder):\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_dim: int,\n",
    "        hidden_dim: int,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.in_dim = in_dim\n",
    "        self.hidden_dim = hidden_dim\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(self.in_dim, self.hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(self.hidden_dim, self.hidden_dim),\n",
    "            nn.ReLU(),\n",
    "        )\n",
    "\n",
    "        # Decode the categories and other attributes separately\n",
    "        self.decoder_categories = nn.Sequential(\n",
    "            nn.Linear(self.hidden_dim, 3),\n",
    "        )\n",
    "\n",
    "        self.decoder_attributes = nn.Sequential(\n",
    "            nn.Linear(self.hidden_dim, 8),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:\n",
    "        out = self.decoder(x)\n",
    "        return [self.decoder_categories(out), self.decoder_attributes(out)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5613963f-a8de-43ce-9319-43de501a03a8",
   "metadata": {},
   "source": [
    "Now let's combine in the DomainModule:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "69671e35-863f-4b31-8955-82fd8d6d7ad4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttributeDomainModule(DomainModule):\n",
    "    def __init__(\n",
    "        self,\n",
    "        latent_dim: int,\n",
    "        hidden_dim: int,\n",
    "        beta: float = 1,\n",
    "        coef_categories: float = 1,\n",
    "        coef_attributes: float = 1,\n",
    "        optim_lr: float = 1e-3,\n",
    "        optim_weight_decay: float = 0,\n",
    "        scheduler_args: SchedulerArgs | None = None,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Defines the Attribute domain module.\n",
    "\n",
    "        Args:\n",
    "            latent_dim (`int`): the latent dimension of the module\n",
    "            hidden_dim (`int`): hidden dimension of the VAE encoders and decoders\n",
    "            beta (`float`): for beta-VAE\n",
    "            coef_categories (`float`): loss coefficient attributed to the category\n",
    "                (Defaults to 1.0)\n",
    "            coef_attributes (`float`): loss coefficient attributed to the rest of the\n",
    "                attributes (Defaults to 1.0)\n",
    "            optim_lr (`float`): learning rate for the optimizer\n",
    "            optim_weight_decay (`float`): weight decay for the optimizer\n",
    "            scheduler_args (`SchedulerArgs | None`): Scheduler arguments\n",
    "        \"\"\"\n",
    "        super().__init__(latent_dim)\n",
    "        self.save_hyperparameters()\n",
    "\n",
    "        self.hidden_dim = hidden_dim\n",
    "        self.coef_categories = coef_categories\n",
    "        self.coef_attributes = coef_attributes\n",
    "\n",
    "        vae_encoder = Encoder(self.hidden_dim, self.latent_dim)\n",
    "        vae_decoder = Decoder(self.latent_dim, self.hidden_dim)\n",
    "        self.vae = VAE(vae_encoder, vae_decoder, beta)\n",
    "\n",
    "        self.optim_lr = optim_lr\n",
    "        self.optim_weight_decay = optim_weight_decay\n",
    "\n",
    "        self.scheduler_args = SchedulerArgs(\n",
    "            max_lr=optim_lr,\n",
    "            total_steps=1,\n",
    "        )\n",
    "        self.scheduler_args.update(scheduler_args or {})\n",
    "\n",
    "    def encode(self, x: Sequence[torch.Tensor]) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Encodes the attributes into the latent representation.\n",
    "\n",
    "        Args:\n",
    "            x (`Sequence[torch.Tensor]`): list with 2 items: the shape category in a one-hot format,\n",
    "                and the attribute vector.\n",
    "        Returns:\n",
    "            `torch.Tensor`: the encoded latent representation\n",
    "        \"\"\"\n",
    "        return self.vae.encode(x)\n",
    "\n",
    "    def decode(self, z: torch.Tensor) -> list[torch.Tensor]:\n",
    "        \"\"\"\n",
    "        Decodes the latent representation to the shape category and attributes.\n",
    "\n",
    "        Args:\n",
    "            z (`torch.Tensor`): the latent representation\n",
    "        Returns:\n",
    "            `list[torch.Tensor]`: list with 2 items: the shape category in a one-hot format,\n",
    "                and the attribute vector.\n",
    "        \"\"\"\n",
    "        return self.vae.decode(z)\n",
    "\n",
    "    def forward(self, x: Sequence[torch.Tensor]) -> list[torch.Tensor]:\n",
    "        return self.decode(self.encode(x))\n",
    "\n",
    "    def compute_loss(self, pred: torch.Tensor, target: torch.Tensor, raw_target: Any) -> LossOutput:\n",
    "        \"\"\"\n",
    "        The loss (in the latent space domain) is simply the MSE between predicted and target latent representations\n",
    "        \"\"\"\n",
    "        return LossOutput(F.mse_loss(pred, target, reduction=\"mean\"))\n",
    "\n",
    "    # Pytorch Lightning functions\n",
    "    # This part is very similar to the visual VAE\n",
    "\n",
    "    def training_step(\n",
    "        self,\n",
    "        batch: Mapping[frozenset[str], Mapping[str, Sequence[torch.Tensor]]],\n",
    "        batch_idx: int,\n",
    "    ) -> torch.Tensor:\n",
    "        x = batch[frozenset([\"attr\"])][\"attr\"]\n",
    "        return self.generic_step(x, \"train\")\n",
    "\n",
    "    def validation_step(self, batch: Mapping[str, Sequence[torch.Tensor]], batch_idx: int) -> torch.Tensor:\n",
    "        x = batch[\"attr\"]\n",
    "        return self.generic_step(x, \"val\")\n",
    "\n",
    "    def generic_step(\n",
    "        self,\n",
    "        x: Sequence[torch.Tensor],\n",
    "        mode: str = \"train\",\n",
    "    ) -> torch.Tensor:\n",
    "        x_categories, x_attributes = x[0], x[1]\n",
    "\n",
    "        (mean, logvar), reconstruction = self.vae(x)\n",
    "        reconstruction_categories = reconstruction[0]\n",
    "        reconstruction_attributes = reconstruction[1]\n",
    "\n",
    "        reconstruction_loss_categories = F.cross_entropy(\n",
    "            reconstruction_categories,\n",
    "            x_categories.argmax(dim=1),\n",
    "            reduction=\"sum\",\n",
    "        )\n",
    "        reconstruction_loss_attributes = gaussian_nll(reconstruction_attributes, torch.tensor(0), x_attributes).sum()\n",
    "\n",
    "        reconstruction_loss = (\n",
    "            self.coef_categories * reconstruction_loss_categories\n",
    "            + self.coef_attributes * reconstruction_loss_attributes\n",
    "        )\n",
    "        kl_loss = kl_divergence_loss(mean, logvar)\n",
    "        total_loss = reconstruction_loss + self.vae.beta * kl_loss\n",
    "\n",
    "        self.log(\n",
    "            f\"{mode}/reconstruction_loss_categories\",\n",
    "            reconstruction_loss_categories,\n",
    "        )\n",
    "        self.log(\n",
    "            f\"{mode}/reconstruction_loss_attributes\",\n",
    "            reconstruction_loss_attributes,\n",
    "        )\n",
    "        self.log(f\"{mode}/reconstruction_loss\", reconstruction_loss)\n",
    "        self.log(f\"{mode}/kl_loss\", kl_loss)\n",
    "        self.log(f\"{mode}/loss\", total_loss)\n",
    "        return total_loss\n",
    "\n",
    "    def configure_optimizers(self) -> dict[str, Any]:\n",
    "        optimizer = torch.optim.AdamW(\n",
    "            self.parameters(),\n",
    "            lr=self.optim_lr,\n",
    "            weight_decay=self.optim_weight_decay,\n",
    "        )\n",
    "        lr_scheduler = OneCycleLR(optimizer, **self.scheduler_args)\n",
    "\n",
    "        return {\n",
    "            \"optimizer\": optimizer,\n",
    "            \"lr_scheduler\": {\n",
    "                \"scheduler\": lr_scheduler,\n",
    "                \"interval\": \"step\",\n",
    "            },\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6499eb8e-3f30-4b2b-9a05-81e383c8bbb7",
   "metadata": {},
   "source": [
    "Now let's instanciate the domain module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e0389cbe-ad0c-4a00-ae46-616fdaefcd09",
   "metadata": {},
   "outputs": [],
   "source": [
    "attr_domain_module = AttributeDomainModule(\n",
    "    latent_dim=config.domain_modules.attribute.latent_dim,\n",
    "    hidden_dim=config.domain_modules.attribute.hidden_dim,\n",
    "    beta=config.domain_modules.attribute.beta,\n",
    "    coef_categories=config.domain_modules.attribute.coef_categories,\n",
    "    coef_attributes=config.domain_modules.attribute.coef_attributes,\n",
    "    optim_lr=config.training.optim.lr,\n",
    "    optim_weight_decay=config.training.optim.weight_decay,\n",
    "    scheduler_args={\n",
    "        \"max_lr\": config.training.optim.max_lr,\n",
    "        \"total_steps\": config.training.max_steps,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d4740e6-2c74-400b-96d2-9e6da7398445",
   "metadata": {},
   "source": [
    "We will use tensorboard to log the losses and reconstructed images. We can use `LogAttributesCallback` from `shimmer_ssd.logging` to\n",
    "log reconstructed images of some image samples.\n",
    "\n",
    "You can update the `train_attr.yaml` config file to change how often images will be updated on tensorboard:\n",
    "```yaml\n",
    "logging:\n",
    "    log_val_medias_every_n_epochs: 1\n",
    "    log_train_medias_every_n_epochs: 1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0e9ae202-1d16-46f1-8c05-c540571bd380",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = TensorBoardLogger(\"logs\", name=\"attr_model\")\n",
    "\n",
    "# Get some image samples to log in tensorboard.\n",
    "val_samples = data_module.get_samples(\"val\", 32)[frozenset([\"attr\"])][\"attr\"]\n",
    "train_samples = data_module.get_samples(\"train\", 32)[frozenset([\"attr\"])][\"attr\"]\n",
    "\n",
    "# Create attr folder where we will save checkpoints\n",
    "(config.default_root_dir / \"attr\").mkdir(exist_ok=True)\n",
    "\n",
    "callbacks: list[Callback] = [\n",
    "    # Will log the validation ground-truth and reconstructions during training\n",
    "    LogAttributesCallback(\n",
    "        val_samples,\n",
    "        log_key=\"images/val_attr\",\n",
    "        mode=\"val\",\n",
    "        every_n_epochs=config.logging.log_val_medias_every_n_epochs,\n",
    "        image_size=32,\n",
    "        ncols=8,\n",
    "    ),\n",
    "    # Will log the training ground-truth and reconstructions during training\n",
    "    LogAttributesCallback(\n",
    "        train_samples,\n",
    "        log_key=\"images/train_attr\",\n",
    "        mode=\"train\",\n",
    "        every_n_epochs=config.logging.log_train_medias_every_n_epochs,\n",
    "        image_size=32,\n",
    "        ncols=8,\n",
    "    ),\n",
    "    # Save the checkpoints\n",
    "    ModelCheckpoint(\n",
    "        dirpath=config.default_root_dir / \"attr\" / f\"version_{logger.version}\",\n",
    "        filename=\"{epoch}\",\n",
    "        monitor=\"val/loss\",\n",
    "        mode=\"min\",\n",
    "        save_last=\"link\",\n",
    "        save_top_k=1,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c78e1465-4ba2-4fc0-b24b-f2693e4f5f5f",
   "metadata": {},
   "source": [
    "For the final model, let's save where the model is saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8670605c-c68c-4ec0-b2b8-e61d9fc9b4b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoints/attr/version_1\n"
     ]
    }
   ],
   "source": [
    "attribute_checkpoint = config.default_root_dir / \"attr\" / f\"version_{logger.version}\"\n",
    "print(attribute_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b5a591-5301-4d87-a5fe-21d0cb60601f",
   "metadata": {},
   "source": [
    "Load tensorboard. You can select the version associated to the previous path. It will appear after the training is started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a5754b7b-791c-49ce-bf73-65ab63db2528",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "aa7e2c56-1d56-46d6-82d0-7ac89598e365",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Reusing TensorBoard on port 6006 (pid 4137456), started 0:01:52 ago. (Use '!kill 4137456' to kill it.)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-e3e70682c2094cac\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-e3e70682c2094cac\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir \"./logs/attr_model\" --reload_interval 30 --reload_task 'auto'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "552083e4-18b5-4a22-af7a-f44e4be2e13f",
   "metadata": {},
   "source": [
    "Let's start the training!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9e4d93a5-c341-4507-a0f9-7695db51afad",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "You are using a CUDA device ('NVIDIA A100 80GB PCIe') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "\n",
      "  | Name | Type | Params | Mode \n",
      "--------------------------------------\n",
      "0 | vae  | VAE  | 11.4 K | train\n",
      "--------------------------------------\n",
      "11.4 K    Trainable params\n",
      "0         Non-trainable params\n",
      "11.4 K    Total params\n",
      "0.046     Total estimated model params size (MB)\n",
      "22        Modules in train mode\n",
      "0         Modules in eval mode\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8439d4271a864673a66c86f9d775598e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |                                                                                            …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc0617cf673647419a799e0283456c58",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |                                                                                                   …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1be98586c9e846dbabdf17f345ec2408",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a3a2e58b9574422bdf3e9ffee0ddf6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4cd98aa8908447c4b01a3f6f36212eb6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0bcb69e39f434ff18c457921b743bd74",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a29642086754adea465bca8db1e56c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f691d1cefdd046e4b8ef0ae0ac1c77b9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7659f0dacc074845b002a3a8d4e55afc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "704d2e28a697428bbc84e89109850dea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_steps=2000` reached.\n",
      "Restoring states from the checkpoint path at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/attr/version_1/epoch=7.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "Loaded model weights from the checkpoint at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/attr/version_1/epoch=7.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "145e67e7edbc43faab03fe7682db9c25",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">          Validate metric           </span>┃<span style=\"font-weight: bold\">            DataLoader 0            </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">            val/kl_loss             </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          16367.7392578125          </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">              val/loss              </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          16316.0771484375          </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">      val/reconstruction_loss       </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          15497.6904296875          </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\"> val/reconstruction_loss_attributes </span>│<span style=\"color: #800080; text-decoration-color: #800080\">           15491.32421875           </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\"> val/reconstruction_loss_categories </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         6.365734577178955          </span>│\n",
       "└────────────────────────────────────┴────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m         Validate metric          \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m           DataLoader 0           \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m           val/kl_loss            \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         16367.7392578125         \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m             val/loss             \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         16316.0771484375         \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m     val/reconstruction_loss      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         15497.6904296875         \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36mval/reconstruction_loss_attributes\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m          15491.32421875          \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36mval/reconstruction_loss_categories\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        6.365734577178955         \u001b[0m\u001b[35m \u001b[0m│\n",
       "└────────────────────────────────────┴────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'val/reconstruction_loss_categories': 6.365734577178955,\n",
       "  'val/reconstruction_loss_attributes': 15491.32421875,\n",
       "  'val/reconstruction_loss': 15497.6904296875,\n",
       "  'val/kl_loss': 16367.7392578125,\n",
       "  'val/loss': 16316.0771484375}]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    logger=logger,\n",
    "    fast_dev_run=config.training.fast_dev_run,\n",
    "    max_steps=config.training.max_steps,\n",
    "    enable_progress_bar=config.training.enable_progress_bar,\n",
    "    default_root_dir=config.default_root_dir,\n",
    "    callbacks=callbacks,\n",
    "    precision=config.training.precision,\n",
    "    accelerator=config.training.accelerator,\n",
    "    devices=config.training.devices,\n",
    ")\n",
    "\n",
    "trainer.fit(attr_domain_module, data_module)\n",
    "trainer.validate(attr_domain_module, data_module, \"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17937142-b242-4cce-8122-372c8dc88baf",
   "metadata": {},
   "source": [
    "## Train a Global Workspace\n",
    "\n",
    "Now that we trained our two unimodal modules, we will train the global workspace. For this training, we will use half of the paired 500,000 samples.\n",
    "To this extent, we need to create a split in the dataset. A dataset split depends on a seed and the proportion of each group of domain.\n",
    "We only need to generate this split once.\n",
    "\n",
    "This can be done with the `shapesd alignment add` command. It needs the following arguments:\n",
    "- `--dataset_path \"DATASET_PATH\"`: the location where the dataset is stored\n",
    "- `--seed SEED` the split seed\n",
    "- `--domain_alignment DOMAIN_1,DOMAIN_2,...DOMAIN_N PROP` the proportion for each domain group. This corresponds to what has been defined in `domain_proportion`\n",
    "\n",
    "When running this command, it will create a file containing the indices of the items available in the train set (update so that it matches what we set in the config file)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ea1d63eb-c931-484f-81da-58edfc090122",
   "metadata": {},
   "outputs": [],
   "source": [
    "!shapesd alignment add --dataset_path \"simple_shapes_dataset\" --seed 0 --domain_alignment attr 1.0 --domain_alignment v 1.0 --domain_alignment attr,v 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "565b121c-f7a5-4bd3-a279-22616387b19a",
   "metadata": {},
   "source": [
    "This time, we will load the config from the extra file `train_gw.yaml` \n",
    "\n",
    "First, let's update `main.yaml` to use the same alignment split:\n",
    "```yaml\n",
    "domain_proportions: \n",
    "    -   domains: [\"v\"]  # unimodal visual passes use 100% of the available data \n",
    "        proportion: 1.0\n",
    "    -   domains: [\"attr\"]  # unimodal attr passes use 100% of the available data\n",
    "        proportion: 1.0\n",
    "    -   domains: [\"v\", \"attr\"]  # paired passes uses 50% of the available data\n",
    "        proportion: 0.5\n",
    "```\n",
    "\n",
    "let's change the selected domains:\n",
    "\n",
    "```yaml\n",
    "domains:\n",
    "    - checkpoint_path: \"./checkpoints/visual/version_0/last.ckpt\"  # update to the actual version\n",
    "      domain_type: v_latents\n",
    "    - checkpoint_path: \"./checkpoints/attr/version_0/last.ckpt\"  # update to the actual version\n",
    "      domain_type: attr\n",
    "```\n",
    "\n",
    "and let's define the global workspace dimenison to 12:\n",
    "```yaml\n",
    "global_workspace:\n",
    "    latent_dim: 12  \n",
    "    \n",
    "    loss_coefficients:\n",
    "        cycles: 1.0\n",
    "        contrastives: 0.1\n",
    "        demi_cycles: 1.0\n",
    "        translations: 1.0\n",
    "\n",
    "    encoders:\n",
    "        hidden_dim: 32\n",
    "        n_layers: 3\n",
    "\n",
    "    decoders:\n",
    "        hidden_dim: 32\n",
    "        n_layers: 3\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb9d8076-115a-4182-bbb8-2599eb49feef",
   "metadata": {},
   "source": [
    "Finally, let's load the config:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0fde1645-60f5-445a-bb50-db1a97de23eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = load_config(\"./config\", use_cli=False, load_files=[\"train_gw.yaml\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72288032-8119-47eb-8721-02490b88152d",
   "metadata": {},
   "source": [
    "Skip the following cell if you have trained the unimodal module yourself. The next cell setups pretrained modules."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0eaa95a-8f85-4cbf-9c12-88911f720a7c",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Run this if you did't train the modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "02e5f922-eb18-4713-a6b6-f2bbda61d5e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading in checkpoints.\n",
      "  0%|                                                | 0.00/288M [00:00<?, ?B/s]\n",
      "  0%|                                        | 254k/288M [00:00<02:05, 2.28MB/s]\n",
      "  2%|▋                                      | 4.64M/288M [00:00<00:11, 25.6MB/s]\n",
      "  5%|█▉                                     | 14.4M/288M [00:00<00:04, 57.7MB/s]\n",
      "  9%|███▍                                   | 25.4M/288M [00:00<00:03, 78.0MB/s]\n",
      " 13%|█████                                  | 37.6M/288M [00:00<00:02, 93.8MB/s]\n",
      " 17%|██████▉                                 | 49.6M/288M [00:00<00:02, 102MB/s]\n",
      " 21%|████████                               | 59.9M/288M [00:00<00:04, 54.7MB/s]\n",
      " 25%|█████████▋                             | 71.4M/288M [00:01<00:03, 66.7MB/s]\n",
      " 29%|███████████▏                           | 82.5M/288M [00:01<00:02, 76.6MB/s]\n",
      " 33%|████████████▊                          | 94.7M/288M [00:01<00:02, 86.8MB/s]\n",
      " 37%|██████████████▊                         | 107M/288M [00:01<00:01, 96.0MB/s]\n",
      " 41%|████████████████▉                        | 119M/288M [00:01<00:01, 103MB/s]\n",
      " 45%|██████████████████▋                      | 131M/288M [00:01<00:01, 107MB/s]\n",
      " 49%|███████████████████▊                    | 142M/288M [00:01<00:02, 63.7MB/s]\n",
      " 53%|█████████████████████▎                  | 153M/288M [00:02<00:01, 72.4MB/s]\n",
      " 57%|██████████████████████▉                 | 165M/288M [00:02<00:01, 81.1MB/s]\n",
      " 61%|████████████████████████▌               | 177M/288M [00:02<00:01, 90.8MB/s]\n",
      " 66%|██████████████████████████▎             | 189M/288M [00:02<00:00, 99.0MB/s]\n",
      " 70%|████████████████████████████▋            | 201M/288M [00:02<00:00, 104MB/s]\n",
      " 74%|██████████████████████████████▎          | 213M/288M [00:02<00:00, 107MB/s]\n",
      " 78%|███████████████████████████████▉         | 224M/288M [00:02<00:00, 107MB/s]\n",
      " 82%|████████████████████████████████▋       | 235M/288M [00:02<00:00, 66.0MB/s]\n",
      " 86%|██████████████████████████████████▍     | 248M/288M [00:03<00:00, 77.3MB/s]\n",
      " 90%|████████████████████████████████████    | 260M/288M [00:03<00:00, 86.6MB/s]\n",
      " 94%|█████████████████████████████████████▋  | 271M/288M [00:03<00:00, 93.6MB/s]\n",
      " 98%|███████████████████████████████████████▍| 283M/288M [00:03<00:00, 99.9MB/s]\n",
      "35108it [00:03, 10295.91it/s]\u001b[A\n",
      "100%|████████████████████████████████████████| 288M/288M [00:03<00:00, 84.3MB/s]\n",
      "Extracting archive...\n",
      "Latent file already exists. Skipping\n"
     ]
    }
   ],
   "source": [
    "# Download checkpoints\n",
    "!ssd download checkpoints\n",
    "!mv checkpoints/checkpoints/* checkpoints/\n",
    "!rm -rf checkpoints/checkpoints\n",
    "\n",
    "# Extract visual latent from pretrained visual domain\n",
    "!ssd extract v \"checkpoints/domain_v.ckpt\" -p \"simple_shapes_dataset\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c21fe165-5f6c-470b-9a6e-ec17789a2799",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update the config\n",
    "checkpoint_path = Path(\"./checkpoints\")\n",
    "\n",
    "config.domain_proportions = {\n",
    "    frozenset([\"v\"]): 1.0,\n",
    "    frozenset([\"attr\"]): 1.0,\n",
    "    frozenset([\"v\", \"attr\"]): 0.5,\n",
    "}\n",
    "\n",
    "config.domains = [\n",
    "    LoadedDomainConfig(\n",
    "        domain_type=DomainModuleVariant.v_latents,\n",
    "        checkpoint_path=checkpoint_path / \"domain_v.ckpt\",\n",
    "    ),\n",
    "    LoadedDomainConfig(\n",
    "        domain_type=DomainModuleVariant.attr,\n",
    "        checkpoint_path=checkpoint_path / \"domain_attr.ckpt\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "config.domain_data_args[\"v_latents\"][\"presaved_path\"] = \"domain_v.npy\"\n",
    "config.global_workspace.latent_dim = 12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14e86345-5644-4007-87fe-5f7a143f4d41",
   "metadata": {},
   "source": [
    "### Load the domains and train\n",
    "We can now load the pretrained unimodal modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e1135c58-53b1-4c8f-9b76-373128d7f218",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we load the pretrained domain modules and define the associated GW encoders and decoders\n",
    "domain_modules, gw_encoders, gw_decoders = load_pretrained_domains(\n",
    "    config.domains,\n",
    "    config.global_workspace.latent_dim,\n",
    "    config.global_workspace.encoders.hidden_dim,\n",
    "    config.global_workspace.encoders.n_layers,\n",
    "    config.global_workspace.decoders.hidden_dim,\n",
    "    config.global_workspace.decoders.n_layers,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebe7e3df-37c9-401b-a481-38609c18ceff",
   "metadata": {},
   "source": [
    "Instanciate the global Workspace class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3c26a008-3d9e-4676-bfb1-5364bbe8e752",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_scheduler(optimizer: Optimizer) -> OneCycleLR:\n",
    "    return OneCycleLR(optimizer, config.training.optim.max_lr, config.training.max_steps)\n",
    "\n",
    "\n",
    "global_workspace = GlobalWorkspace2Domains(\n",
    "    domain_modules,\n",
    "    gw_encoders,\n",
    "    gw_decoders,\n",
    "    config.global_workspace.latent_dim,\n",
    "    config.global_workspace.loss_coefficients,\n",
    "    config.training.optim.lr,\n",
    "    config.training.optim.weight_decay,\n",
    "    scheduler=get_scheduler,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9171f84d-e76e-44e7-95c7-246f064e664d",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_classes = get_default_domains([\"v_latents\", \"attr\"])\n",
    "\n",
    "data_module = SimpleShapesDataModule(\n",
    "    config.dataset.path,\n",
    "    domain_classes,\n",
    "    config.domain_proportions,\n",
    "    batch_size=config.training.batch_size,\n",
    "    num_workers=config.training.num_workers,\n",
    "    seed=config.seed,\n",
    "    domain_args=config.domain_data_args,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf92cee-a17a-4db2-960b-91e375235de5",
   "metadata": {},
   "source": [
    "Add a tensorboard logger to follow the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "145c0039-b27a-4ab3-a262-abb87a00d152",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = TensorBoardLogger(\"logs\", name=\"gw\")\n",
    "\n",
    "# Get some image samples to log in tensorboard.\n",
    "train_samples = data_module.get_samples(\"train\", 32)\n",
    "val_samples = data_module.get_samples(\"val\", 32)\n",
    "\n",
    "# split the unique group in validation into individual groups for logging\n",
    "for domains in val_samples:\n",
    "    for domain in domains:\n",
    "        val_samples[frozenset([domain])] = {domain: val_samples[domains][domain]}\n",
    "    break\n",
    "# Create attr folder where we will save checkpoints\n",
    "(config.default_root_dir / \"gw\").mkdir(exist_ok=True)\n",
    "\n",
    "callbacks: list[Callback] = [\n",
    "    # Will log the validation ground-truth and reconstructions during training\n",
    "    LogGWImagesCallback(\n",
    "        val_samples,\n",
    "        log_key=\"images/val\",\n",
    "        mode=\"val\",\n",
    "        every_n_epochs=config.logging.log_val_medias_every_n_epochs,\n",
    "        filter=config.logging.filter_images,\n",
    "    ),\n",
    "    # Will log the training ground-truth and reconstructions during training\n",
    "    LogGWImagesCallback(\n",
    "        train_samples,\n",
    "        log_key=\"images/train\",\n",
    "        mode=\"train\",\n",
    "        every_n_epochs=config.logging.log_train_medias_every_n_epochs,\n",
    "        filter=config.logging.filter_images,\n",
    "    ),\n",
    "    # Save the checkpoints\n",
    "    ModelCheckpoint(\n",
    "        dirpath=config.default_root_dir / \"gw\" / f\"version_{logger.version}\",\n",
    "        filename=\"{epoch}\",\n",
    "        monitor=\"val/loss\",\n",
    "        mode=\"min\",\n",
    "        save_last=\"link\",\n",
    "        save_top_k=1,\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59ae4778-9615-4415-abad-420787e66d2c",
   "metadata": {},
   "source": [
    "For the final model, let's save where the model is saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b98b6642-ae5c-4cf8-9600-07882cd9d4ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "checkpoints/gw/version_2\n"
     ]
    }
   ],
   "source": [
    "gw_checkpoint = config.default_root_dir / \"gw\" / f\"version_{logger.version}\"\n",
    "print(gw_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3cd34b0-60d7-4899-984f-3a5b7ccef6f0",
   "metadata": {},
   "source": [
    "Load tensorboard. You can select the version associated to the previous path. It will appear after the training is started."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0926b71e-1079-4108-a462-7f43546e77b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b12eb7b7-1459-406d-aaf5-5b638dae27ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "      <iframe id=\"tensorboard-frame-4be75f3e06dea1aa\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
       "      </iframe>\n",
       "      <script>\n",
       "        (function() {\n",
       "          const frame = document.getElementById(\"tensorboard-frame-4be75f3e06dea1aa\");\n",
       "          const url = new URL(\"/\", window.location);\n",
       "          const port = 6006;\n",
       "          if (port) {\n",
       "            url.port = port;\n",
       "          }\n",
       "          frame.src = url;\n",
       "        })();\n",
       "      </script>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir \"./logs/gw\" --reload_interval 30 --reload_task 'auto'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82764468-75c4-4403-89e4-269ac0ae8779",
   "metadata": {},
   "source": [
    "And train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "cf6de9e3-d4cd-4246-94df-4a677f8866b0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "\n",
      "  | Name          | Type                  | Params | Mode \n",
      "----------------------------------------------------------------\n",
      "0 | gw_mod        | GWModule              | 5.9 M  | train\n",
      "1 | selection_mod | SingleDomainSelection | 0      | train\n",
      "2 | loss_mod      | GWLosses2Domains      | 5.9 M  | train\n",
      "----------------------------------------------------------------\n",
      "15.9 K    Trainable params\n",
      "5.8 M     Non-trainable params\n",
      "5.9 M     Total params\n",
      "23.428    Total estimated model params size (MB)\n",
      "47        Modules in train mode\n",
      "57        Modules in eval mode\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7907157ed2404360bf6613ae4ec30311",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |                                                                                            …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb1c358b657c455a858e60ece12ec8df",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |                                                                                                   …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ded8b395310e45e4af1c5c431ce9b690",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7090a88b7bef4a8eae426b6f75b40d53",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "168c9a4c1b0947f380aa144870b73ff2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f6211c008094ca6a40317512fe61a2f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef158f26eeb64474b6e51ba63bfcaef1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "13e689c08e394397af715204465f59b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d962c807781f4395bd6bf2b8f347086d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a0644a5109824e8e83c1ee8036ca514b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52bc46a04be341ed8a0d15cc14bed750",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6c05a13b84b46898ae335e20ee770d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "60716e33b618458580f2598a2bc8c791",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "344f12f6e53543fbb3387a6fb97d4050",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c133bbdc6a4e4beda39f3f8cc9da445c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6db128282910453d9a4d2619c2735dfa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0155604002a4c6a82cd6a956cb69c32",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf332f64b7e440d1879afdbe13ec541f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9aae6d0b529e444dad67175751482134",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a868fc8fdf1046d99ed0e1bfe84972e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0bac39da29ef493aac4775a930660ae2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2af3de4d9e2c4285878b446b3dc32712",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db30c0cbe8d14daebc8e4f0544dfb449",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9771ca8e88054908b1816f3e5b1239a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6bc9b0e9f6a4b2f89993a1eb7c146a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdb595c0656c4bc690871dc06593f042",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b51ee9eb61994ed1b9a21bdd22bc558a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "03cc2149c4cc41518ac143da639fbf46",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "790313dd30b84f1f9bb64eaf6cac3332",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fb76922e4a2f4f5498a994f91ca9ac2e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2109394bd6c54dd09d15193071b060e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3d40700f5fd404397d5ae5f3062e67e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73eeb98b5b824d84a839a8b900c44188",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8a7ac49f2ba3448e9e44217cd9750fb1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_steps=8000` reached.\n",
      "Restoring states from the checkpoint path at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/gw/version_2/epoch=31.ckpt\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]\n",
      "Loaded model weights from the checkpoint at /home/bdevillers/projects/shimmer-ssd-tutorials/checkpoints/gw/version_2/epoch=31.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "525644d135df4a119f87afc7502ef1dd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |                                                                                                 …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\">                Validate metric                 </span>┃<span style=\"font-weight: bold\">                  DataLoader 0                  </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">       val/contrastive_v_latents_and_attr       </span>│<span style=\"color: #800080; text-decoration-color: #800080\">               1.4367294311523438               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\"> val/contrastive_v_latents_and_attr_logit_scale </span>│<span style=\"color: #800080; text-decoration-color: #800080\">               14.285717964172363               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">                val/contrastives                </span>│<span style=\"color: #800080; text-decoration-color: #800080\">               1.4367294311523438               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">        val/cycle_attr_through_v_latents        </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.11297674477100372               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">        val/cycle_v_latents_through_attr        </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.21616032719612122               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">                   val/cycles                   </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.16456852853298187               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">              val/demi_cycle_attr               </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.03258915990591049               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">            val/demi_cycle_v_latents            </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.16235917806625366               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">                val/demi_cycles                 </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.09747415781021118               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">                    val/loss                    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.15514861047267914               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">       val/translation_attr_to_v_latents        </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.25512421131134033               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">       val/translation_v_latents_to_attr        </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.17463333904743195               </span>│\n",
       "│<span style=\"color: #008080; text-decoration-color: #008080\">                val/translations                </span>│<span style=\"color: #800080; text-decoration-color: #800080\">              0.21487878262996674               </span>│\n",
       "└────────────────────────────────────────────────┴────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1m               Validate metric                \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                 DataLoader 0                 \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│\u001b[36m \u001b[0m\u001b[36m      val/contrastive_v_latents_and_attr      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m              1.4367294311523438              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36mval/contrastive_v_latents_and_attr_logit_scale\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m              14.285717964172363              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m               val/contrastives               \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m              1.4367294311523438              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m       val/cycle_attr_through_v_latents       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.11297674477100372              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m       val/cycle_v_latents_through_attr       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.21616032719612122              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m                  val/cycles                  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.16456852853298187              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m             val/demi_cycle_attr              \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.03258915990591049              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m           val/demi_cycle_v_latents           \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.16235917806625366              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m               val/demi_cycles                \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.09747415781021118              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m                   val/loss                   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.15514861047267914              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m      val/translation_attr_to_v_latents       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.25512421131134033              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m      val/translation_v_latents_to_attr       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.17463333904743195              \u001b[0m\u001b[35m \u001b[0m│\n",
       "│\u001b[36m \u001b[0m\u001b[36m               val/translations               \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m             0.21487878262996674              \u001b[0m\u001b[35m \u001b[0m│\n",
       "└────────────────────────────────────────────────┴────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "[{'val/demi_cycle_v_latents': 0.16235917806625366,\n",
       "  'val/demi_cycle_attr': 0.03258915990591049,\n",
       "  'val/demi_cycles': 0.09747415781021118,\n",
       "  'val/cycle_v_latents_through_attr': 0.21616032719612122,\n",
       "  'val/cycle_attr_through_v_latents': 0.11297674477100372,\n",
       "  'val/cycles': 0.16456852853298187,\n",
       "  'val/translation_attr_to_v_latents': 0.25512421131134033,\n",
       "  'val/translation_v_latents_to_attr': 0.17463333904743195,\n",
       "  'val/translations': 0.21487878262996674,\n",
       "  'val/contrastive_v_latents_and_attr': 1.4367294311523438,\n",
       "  'val/contrastives': 1.4367294311523438,\n",
       "  'val/contrastive_v_latents_and_attr_logit_scale': 14.285717964172363,\n",
       "  'val/loss': 0.15514861047267914}]"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    logger=logger,\n",
    "    max_steps=config.training.max_steps,\n",
    "    default_root_dir=config.default_root_dir,\n",
    "    callbacks=callbacks,\n",
    "    precision=config.training.precision,\n",
    "    accelerator=config.training.accelerator,\n",
    "    devices=config.training.devices,\n",
    ")\n",
    "\n",
    "trainer.fit(global_workspace, data_module)\n",
    "trainer.validate(global_workspace, data_module, \"best\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f99d2f8-e3d4-4e42-bda0-bc949bf1d248",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Run this if you did't train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "baeef54d-ab50-490d-931b-8fa8d71e9026",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update the config\n",
    "checkpoint_path = Path(\"./checkpoints\")\n",
    "\n",
    "config.domain_proportions = {\n",
    "    frozenset([\"v\"]): 1.0,\n",
    "    frozenset([\"attr\"]): 1.0,\n",
    "    frozenset([\"v\", \"attr\"]): 0.5,\n",
    "}\n",
    "\n",
    "config.domains = [\n",
    "    LoadedDomainConfig(\n",
    "        domain_type=DomainModuleVariant.v_latents,\n",
    "        checkpoint_path=checkpoint_path / \"domain_v.ckpt\",\n",
    "    ),\n",
    "    LoadedDomainConfig(\n",
    "        domain_type=DomainModuleVariant.attr,\n",
    "        checkpoint_path=checkpoint_path / \"domain_attr.ckpt\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "config.domain_data_args[\"v_latents\"][\"presaved_path\"] = \"domain_v.npy\"\n",
    "config.global_workspace.latent_dim = 12\n",
    "# And now we load the GW checkpoint\n",
    "checkpoint_path = Path(\"./checkpoints\")\n",
    "checkpoint = checkpoint_path / \"gw-attr-v-half-paired-data.ckpt\"\n",
    "\n",
    "# we load the pretrained domain modules and define the associated GW encoders and decoders\n",
    "domain_modules, gw_encoders, gw_decoders = load_pretrained_domains(\n",
    "    config.domains,\n",
    "    config.global_workspace.latent_dim,\n",
    "    config.global_workspace.encoders.hidden_dim,\n",
    "    config.global_workspace.encoders.n_layers,\n",
    "    config.global_workspace.decoders.hidden_dim,\n",
    "    config.global_workspace.decoders.n_layers,\n",
    ")\n",
    "\n",
    "global_workspace = GlobalWorkspace2Domains.load_from_checkpoint(\n",
    "    checkpoint,\n",
    "    domain_mods=domain_modules,\n",
    "    gw_encoders=gw_encoders,\n",
    "    gw_decoders=gw_decoders,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aefbfff6-8cbc-4a72-8e34-47c2d28f0e6d",
   "metadata": {},
   "source": [
    "## Play with the global workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e0cc65be-5456-4f3a-8216-0d061deec60e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import math\n",
    "\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from ipywidgets import interact, interact_manual\n",
    "from PIL import Image\n",
    "from shimmer_ssd.logging import attribute_image_grid\n",
    "from torch.nn.functional import one_hot\n",
    "\n",
    "from simple_shapes_dataset.cli import generate_image\n",
    "%matplotlib widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b204af9d-df27-4462-9dfc-b45d8e7eb3d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4f37e3214cb4738951b5144b9192720",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(Dropdown(description='cat', options=('Triangle', 'Egg', 'Diamond'), value='Triangle'), F…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "global_workspace.to(device)\n",
    "\n",
    "cat2idx = {\"Diamond\": 0, \"Egg\": 1, \"Triangle\": 2}\n",
    "\n",
    "\n",
    "def get_image(cat, x, y, size, rot, color_r, color_g, color_b):\n",
    "    fig, ax = plt.subplots(figsize=(32, 32), dpi=1)\n",
    "    # The dataset generatoion tool has function to generate a matplotlib shape\n",
    "    # from the attributes. \n",
    "    generate_image(\n",
    "        ax,\n",
    "        cat2idx[cat],\n",
    "        [int(x * 18 + 7), int(y * 18 + 7)],\n",
    "        size * 7 + 7,\n",
    "        rot * 2 * math.pi,\n",
    "        np.array([color_r * 255, color_g * 255, color_b * 255]),\n",
    "        imsize=32,\n",
    "    )\n",
    "    ax.set_facecolor(\"black\")\n",
    "    plt.tight_layout(pad=0)\n",
    "    # Return this as a PIL Image.\n",
    "    # This is to have the same dpi as saved images\n",
    "    # otherwise matplotlib will render this in very high quality\n",
    "    buf = io.BytesIO()\n",
    "    fig.savefig(buf)\n",
    "    buf.seek(0)\n",
    "    image = Image.open(buf)\n",
    "    plt.close(fig)\n",
    "    return image\n",
    "\n",
    "\n",
    "@interact(\n",
    "    cat=[\"Triangle\", \"Egg\", \"Diamond\"],\n",
    "    x=(0, 1, 0.1),\n",
    "    y=(0, 1, 0.1),\n",
    "    rot=(0, 1, 0.1),\n",
    "    size=(0, 1, 0.1),\n",
    "    color_r=(0, 1, 0.1),\n",
    "    color_g=(0, 1, 0.1),\n",
    "    color_b=(0, 1, 0.1),\n",
    ")\n",
    "def play_with_gw(\n",
    "    cat: str = \"Triangle\",\n",
    "    x: float = 0.5,\n",
    "    y: float = 0.5,\n",
    "    rot: float = 0.5,\n",
    "    size: float = 0.5,\n",
    "    color_r: float = 1,\n",
    "    color_g: float = 0,\n",
    "    color_b: float = 0,\n",
    "):\n",
    "    fig, axes = plt.subplots(1, 2)\n",
    "    image = get_image(cat, x, y, size, rot, color_r, color_g, color_b)\n",
    "    axes[0].set_facecolor(\"black\")\n",
    "    axes[0].set_title(\"Original image from attributes\")\n",
    "    axes[0].set_xticks([])\n",
    "    axes[0].set_yticks([])\n",
    "    axes[0].imshow(image)\n",
    "\n",
    "    # normalize the attribute for the global workspace.\n",
    "    category = one_hot(torch.tensor([cat2idx[cat]]), 3)\n",
    "    rotx = math.cos(rot * 2 * math.pi)\n",
    "    roty = math.sin(rot * 2 * math.pi)\n",
    "    attributes = torch.tensor(\n",
    "        [[x * 2 - 1, y * 2 - 1, size * 2 - 1, rotx, roty, color_r * 2 - 1, color_g * 2 - 1, color_b * 2 - 1]]\n",
    "    )\n",
    "    samples = [category.to(device), attributes.to(device)]\n",
    "    attr_gw_latent = global_workspace.gw_mod.encode({\"attr\": global_workspace.encode_domain(samples, \"attr\")})\n",
    "    gw_latent = global_workspace.gw_mod.fuse(\n",
    "        attr_gw_latent, {\"attr\": torch.ones(attr_gw_latent[\"attr\"].size(0)).to(device)}\n",
    "    )\n",
    "    decoded_latents = global_workspace.gw_mod.decode(gw_latent)[\"v_latents\"]\n",
    "    decoded_images = (\n",
    "        global_workspace.domain_mods[\"v_latents\"]\n",
    "        .decode_images(decoded_latents)[0]\n",
    "        .permute(1, 2, 0)\n",
    "        .detach()\n",
    "        .cpu()\n",
    "        .numpy()\n",
    "    )\n",
    "    axes[1].imshow(decoded_images)\n",
    "    axes[1].set_xticks([])\n",
    "    axes[1].set_yticks([])\n",
    "    axes[1].set_title(\"Translated image through GW\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f83fca4f-0731-439a-ae8c-34d6c6d80474",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "473ffe9e-0518-45ae-9a52-06090722d711",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}