create script for dry-running all example tasks

Author: henrykrumb

ncalab/models/applications/classification/classificationNCAhead.py

@@ -34,8 +34,8 @@ def forward(self, x):
         return self.classifier(x)
 
     def freeze(self, freeze_last: bool = False):
-        layers = self.classifier
+        layers = [L for L in self.classifier.modules()]
         if not freeze_last:
-            layers = self.classifier[:-1]
+            layers = layers[:-1]
         for layer in layers:
             layer.requires_grad_(False)

ncalab/models/wrappers/cascadeNCA.py

@@ -1,4 +1,5 @@
-from typing import Dict, Optional, Tuple, List
+import copy
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch  # type: ignore[import-untyped]
@@ -115,11 +116,11 @@ def __init__(
         self.steps = steps
 
         self.single_model = single_model
-        self.models: nn.ModuleList | List[BasicNCAModel]
+        self.models: List[BasicNCAModel]
         if single_model:
             self.models = [wrapped for _ in scales]
         else:
-            self.models = nn.ModuleList([wrapped for _ in scales])
+            self.models = [copy.deepcopy(wrapped) for _ in scales]
 
     def forward(self, x: torch.Tensor, *args, **kwargs) -> Prediction:
         """

ncalab/training/trainer.py

@@ -28,7 +28,7 @@ class BasicNCATrainer:
     def __init__(
         self,
         nca: BasicNCAModel,
-        model_path: Path | PosixPath,
+        model_path: Optional[Path | PosixPath],
         gradient_clipping: bool = False,
         lr: Optional[float] = None,
         lr_gamma: float = 0.99,

scripts/dry_train_examples.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+set -e
+
+uv run tasks/class_cifar10/train_class_cifar10.py --dry --max-epochs 2
+uv run tasks/growing_emoji/train_growing_emoji.py --dry --epochs 2
+uv run tasks/class_medmnist/train_class_dermamnist.py --dry --max-epochs 2
+uv run tasks/class_medmnist/train_class_bloodmnist.py --dry --max-epochs 2
+uv run tasks/class_medmnist/train_class_pathmnist.py --dry --max-epochs 2
+uv run tasks/segmentation_kvasir_seg/train_segmentation_kvasir_seg.py --dry --max-epochs 2
+uv run tasks/selfclass_mnist/train_selfclass_mnist.py --dry --max-epochs 2

scripts/train_examples.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+set -e
 
 uv run tasks/class_cifar10/train_class_cifar10.py
 uv run tasks/growing_emoji/train_growing_emoji.py

tasks/class_cifar10/train_class_cifar10.py

@@ -37,6 +37,8 @@ def train_class_cifar10(
     hidden_channels: int,
     gpu: bool,
     gpu_index: int,
+    max_epochs: int,
+    dry: bool,
 ):
     comment = "CIFAR10"
     comment += f"_hidden_{hidden_channels}"
@@ -45,7 +47,7 @@ def train_class_cifar10(
     comment += f"_AM_{alive_mask}"
     comment += f"_TE_{use_temporal_encoding}"
 
-    writer = SummaryWriter(comment=comment)
+    writer = SummaryWriter(comment=comment) if not dry else None
 
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
 
@@ -141,17 +143,18 @@ def train_class_cifar10(
     # Train the NCA model
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "classification_cifar10",
+        WEIGHTS_PATH / "classification_cifar10" if not dry else None,
         batch_repeat=2,
-        max_epochs=120,
+        max_epochs=max_epochs,
         gradient_clipping=gradient_clipping,
     )
     trainer.train(
         loader_train,
         loader_val,
         summary_writer=writer,
     )
-    writer.close()
+    if writer is not None:
+        writer.close()
 
 
 @click.command()
@@ -163,12 +166,18 @@ def train_class_cifar10(
 @click.option(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
-def main(batch_size, hidden_channels, gpu: bool, gpu_index: int):
+@click.option("--max-epochs", "-E", type=int, default=120)
+@click.option("--dry", "-D", is_flag=True)
+def main(
+    batch_size, hidden_channels, gpu: bool, gpu_index: int, max_epochs: int, dry: bool
+):
     train_class_cifar10(
         batch_size=batch_size,
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,
+        max_epochs=max_epochs,
+        dry=dry,
     )
 
 

tasks/class_medmnist/train_class_bloodmnist.py

@@ -38,6 +38,8 @@ def train_class_bloodmnist(
     hidden_channels: int,
     gpu: bool,
     gpu_index: int,
+    max_epochs: int,
+    dry: bool,
 ):
     print_NCALab_banner()
 
@@ -48,7 +50,7 @@ def train_class_bloodmnist(
     comment += f"_AM_{alive_mask}"
     comment += f"_TE_{use_temporal_encoding}"
 
-    writer = SummaryWriter(comment=comment)
+    writer = SummaryWriter(comment=comment) if not dry else None
 
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
 
@@ -84,22 +86,23 @@ def train_class_bloodmnist(
         pad_noise=pad_noise,
         use_temporal_encoding=use_temporal_encoding,
         class_names=list(INFO["bloodmnist"]["label"].values()),
+        training_timesteps=32,
+        inference_timesteps=32,
     )
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "classification_bloodmnist",
+        WEIGHTS_PATH / "classification_bloodmnist" if not dry else None,
         batch_repeat=2,
-        max_epochs=40,
+        max_epochs=max_epochs,
         gradient_clipping=gradient_clipping,
-        steps_range=(32, 33),
-        steps_validation=32,
     )
     trainer.train(
         loader_train,
         loader_val,
         summary_writer=writer,
     )
-    writer.close()
+    if writer is not None:
+        writer.close()
 
 
 @click.command()
@@ -111,12 +114,18 @@ def train_class_bloodmnist(
 @click.option(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
-def main(batch_size, hidden_channels, gpu: bool, gpu_index: int):
+@click.option("--max-epochs", "-E", type=int, default=40)
+@click.option("--dry", "-D", is_flag=True)
+def main(
+    batch_size, hidden_channels, gpu: bool, gpu_index: int, max_epochs: int, dry: bool
+):
     train_class_bloodmnist(
         batch_size=batch_size,
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,
+        max_epochs=max_epochs,
+        dry=dry,
     )
 
 

tasks/class_medmnist/train_class_dermamnist.py

@@ -38,6 +38,8 @@ def train_class_dermamnist(
     hidden_channels: int,
     gpu: bool,
     gpu_index: int,
+    max_epochs: int,
+    dry: bool,
 ):
     print_NCALab_banner()
 
@@ -48,7 +50,7 @@ def train_class_dermamnist(
     comment += f"_AM_{alive_mask}"
     comment += f"_TE_{use_temporal_encoding}"
 
-    writer = SummaryWriter(comment=comment)
+    writer = SummaryWriter(comment=comment) if not dry else None
 
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
 
@@ -112,23 +114,24 @@ def train_class_dermamnist(
         pad_noise=pad_noise,
         use_temporal_encoding=use_temporal_encoding,
         class_names=list(INFO["dermamnist"]["label"].values()),
+        training_timesteps=32,
+        inference_timesteps=32,
     )
 
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "classification_dermamnist",
+        WEIGHTS_PATH / "classification_dermamnist" if not dry else None,
         batch_repeat=2,
-        max_epochs=40,
+        max_epochs=max_epochs,
         gradient_clipping=gradient_clipping,
-        steps_range=(32, 33),
-        steps_validation=32,
     )
     trainer.train(
         loader_train,
         loader_val,
         summary_writer=writer,
     )
-    writer.close()
+    if writer is not None:
+        writer.close()
 
 
 @click.command()
@@ -140,12 +143,18 @@ def train_class_dermamnist(
 @click.option(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
-def main(batch_size, hidden_channels, gpu: bool, gpu_index: int):
+@click.option("--max-epochs", "-E", type=int, default=40)
+@click.option("--dry", "-D", is_flag=True)
+def main(
+    batch_size, hidden_channels, gpu: bool, gpu_index: int, max_epochs: int, dry: bool
+):
     train_class_dermamnist(
         batch_size=batch_size,
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,
+        max_epochs=max_epochs,
+        dry=dry,
     )
 
 

tasks/class_medmnist/train_class_pathmnist.py

@@ -36,6 +36,8 @@ def train_class_pathmnist(
     hidden_channels: int,
     gpu: bool,
     gpu_index: int,
+    max_epochs: int,
+    dry: bool,
 ):
     comment = "PathMNIST"
     comment += f"_hidden_{hidden_channels}"
@@ -44,7 +46,7 @@ def train_class_pathmnist(
     comment += f"_AM_{alive_mask}"
     comment += f"_TE_{use_temporal_encoding}"
 
-    writer = SummaryWriter(comment=comment)
+    writer = SummaryWriter(comment=comment) if not dry else None
 
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
 
@@ -80,22 +82,23 @@ def train_class_pathmnist(
         pad_noise=pad_noise,
         use_temporal_encoding=use_temporal_encoding,
         class_names=list(INFO["pathmnist"]["label"].values()),
+        training_timesteps=32,
+        inference_timesteps=32,
     )
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "classification_pathmnist",
+        WEIGHTS_PATH / "classification_pathmnist" if not dry else None,
         batch_repeat=2,
-        max_epochs=10,
+        max_epochs=max_epochs,
         gradient_clipping=gradient_clipping,
-        steps_range=(32, 33),
-        steps_validation=32,
     )
     trainer.train(
         loader_train,
         loader_val,
         summary_writer=writer,
     )
-    writer.close()
+    if writer is not None:
+        writer.close()
 
 
 @click.command()
@@ -107,12 +110,18 @@ def train_class_pathmnist(
 @click.option(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
-def main(batch_size, hidden_channels, gpu: bool, gpu_index: int):
+@click.option("--max-epochs", "-E", type=int, default=10)
+@click.option("--dry", "-D", is_flag=True)
+def main(
+    batch_size, hidden_channels, gpu: bool, gpu_index: int, max_epochs: int, dry: bool
+):
     train_class_pathmnist(
         batch_size=batch_size,
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,
+        max_epochs=max_epochs,
+        dry=dry,
     )
 
 

tasks/growing_emoji/train_growing_emoji.py

@@ -28,7 +28,12 @@
 
 
 def train_growing_emoji(
-    batch_size: int, hidden_channels: int, gpu: bool, gpu_index: int, max_epochs: int
+    batch_size: int,
+    hidden_channels: int,
+    gpu: bool,
+    gpu_index: int,
+    max_epochs: int,
+    dry: bool,
 ):
     """
     Main function to run the "growing emoji" example task.
@@ -40,7 +45,7 @@ def train_growing_emoji(
     fix_random_seed()
 
     # Create tensorboard summary writer
-    writer = SummaryWriter(comment=" Growing NCA (Lizard)")
+    writer = SummaryWriter(comment=" Growing NCA (Lizard)") if not dry else None
 
     # Select device, try to use GPU or fall back to CPU
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
@@ -63,12 +68,13 @@ def train_growing_emoji(
     # Create Trainer and run training
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "ncalab_growing_emoji",
+        WEIGHTS_PATH / "ncalab_growing_emoji" if not dry else None,
         max_epochs=max_epochs,
         pool=pool,
     )
     trainer.train(dataloader_train, summary_writer=writer, save_every=500)
-    writer.close()
+    if writer is not None:
+        writer.close()
 
 
 @click.command()
@@ -81,13 +87,15 @@ def train_growing_emoji(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
 @click.option("--epochs", "-e", type=int, default=5000)
-def main(batch_size, hidden_channels, gpu, gpu_index, epochs):
+@click.option("--dry", "-D", is_flag=True)
+def main(batch_size, hidden_channels, gpu, gpu_index: int, epochs: int, dry: bool):
     train_growing_emoji(
         batch_size=batch_size,
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,
         max_epochs=epochs,
+        dry=dry,
     )