medmnist classification tasks: share parameters between training and evaluation

Author: henrykrumb

tasks/class_medmnist/README.md

@@ -0,0 +1,3 @@
+# MedMNIST Classification
+
+This example uses a classification NCA model.

tasks/class_medmnist/eval_class_bloodmnist.py

@@ -5,63 +5,71 @@
 root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
 sys.path.append(root_dir)
 
-from ncalab import ClassificationNCAModel, WEIGHTS_PATH, get_compute_device, pad_input
+from ncalab import ClassificationNCAModel, get_compute_device
 
 import click
 
-from medmnist import PathMNIST  # type: ignore[import-untyped]
+from medmnist import INFO, DermaMNIST  # type: ignore[import-untyped]
 
 import torch  # type: ignore[import-untyped]
 from torchvision import transforms  # type: ignore[import-untyped]
 from torchvision.transforms import v2  # type: ignore[import-untyped]
 
-import numpy as np
-
 import torchmetrics
 import torchmetrics.classification
 
 from tqdm import tqdm
 
+from train_class_dermamnist import (
+    pad_noise,
+    alive_mask,
+    use_temporal_encoding,
+    fire_rate,
+    WEIGHTS_PATH,
+)
 
 T = transforms.Compose(
     [
         v2.ToImage(),
         v2.ToDtype(torch.float, scale=True),
         v2.ConvertImageDtype(dtype=torch.float32),
-        transforms.RandomHorizontalFlip(),
-        transforms.RandomVerticalFlip(),
+        transforms.Normalize((0.5,), (0.225,)),
     ]
 )
 
 
-def eval_selfclass_pathmnist(
+def eval_selfclass_dermamnist(
     hidden_channels: int,
     gpu,
     gpu_index,
 ):
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
 
-    dataset_test = PathMNIST(split="test", download=True, transform=T)
+    dataset_test = DermaMNIST(split="test", download=True, transform=T)
     loader_test = torch.utils.data.DataLoader(
         dataset_test, shuffle=False, batch_size=32
     )
 
-    num_classes = 9
+    num_classes = len(INFO["dermamnist"]["label"])
     nca = ClassificationNCAModel(
         device,
         num_image_channels=3,
         num_hidden_channels=hidden_channels,
         hidden_size=128,
         num_classes=num_classes,
-        use_alive_mask=False,
-        fire_rate=0.5,
+        use_alive_mask=alive_mask,
+        fire_rate=fire_rate,
+        use_temporal_encoding=use_temporal_encoding,
+        pad_noise=pad_noise,
     )
     nca.load_state_dict(
-        torch.load(WEIGHTS_PATH / "selfclass_pathmnist.best.pth", weights_only=True)
+        torch.load(
+            WEIGHTS_PATH / "classification_dermamnist" / "best_model.pth",
+            weights_only=True,
+        )
     )
 
-    model_parameters = filter(lambda p: p.requires_grad, nca.parameters())
-    params = sum([np.prod(p.size()) for p in model_parameters])
+    params = nca.num_trainable_parameters()
     print(f"Trainable parameters: {params}")
     print(f"That is {4 * params / 1000} kB")
 
@@ -70,24 +78,24 @@ def eval_selfclass_pathmnist(
 
     macro_acc = torchmetrics.classification.MulticlassAccuracy(
         average="macro", num_classes=num_classes
-    )
+    ).to(device)
     micro_acc = torchmetrics.classification.MulticlassAccuracy(
         average="micro", num_classes=num_classes
-    )
+    ).to(device)
     macro_auc = torchmetrics.classification.MulticlassAUROC(
         average="macro",
         num_classes=num_classes,
-    )
+    ).to(device)
     micro_f1 = torchmetrics.classification.MulticlassF1Score(
         average="micro", num_classes=num_classes
-    )
+    ).to(device)
     for sample in tqdm(iter(loader_test)):
         x, y = sample
-        x = pad_input(x, nca, noise=True)
         x = x.float().to(device)
         steps = 72
         y_prob = nca.classify(x, steps, reduce=False)
         y = y.squeeze().to(device)
+
         macro_acc.update(y_prob, y)
         micro_acc.update(y_prob, y)
         macro_auc.update(y_prob, y)
@@ -104,15 +112,15 @@ def eval_selfclass_pathmnist(
 
 
 @click.command()
-@click.option("--hidden-channels", "-H", default=20, type=int)
+@click.option("--hidden-channels", "-H", default=30, type=int)
 @click.option(
     "--gpu/--no-gpu", is_flag=True, default=True, help="Try using the GPU if available."
 )
 @click.option(
     "--gpu-index", type=int, default=0, help="Index of GPU to use, if --gpu in use."
 )
 def main(hidden_channels, gpu: bool, gpu_index: int):
-    eval_selfclass_pathmnist(
+    eval_selfclass_dermamnist(
         hidden_channels=hidden_channels,
         gpu=gpu,
         gpu_index=gpu_index,

tasks/class_medmnist/eval_class_dermamnist.py

@@ -5,31 +5,35 @@
 root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
 sys.path.append(root_dir)
 
-from ncalab import ClassificationNCAModel, WEIGHTS_PATH, get_compute_device, pad_input
+from ncalab import ClassificationNCAModel, get_compute_device
 
 import click
 
-from medmnist import PathMNIST  # type: ignore[import-untyped]
+from medmnist import INFO, PathMNIST  # type: ignore[import-untyped]
 
 import torch  # type: ignore[import-untyped]
 from torchvision import transforms  # type: ignore[import-untyped]
 from torchvision.transforms import v2  # type: ignore[import-untyped]
 
-import numpy as np
-
 import torchmetrics
 import torchmetrics.classification
 
 from tqdm import tqdm
 
+from train_class_pathmnist import (
+    pad_noise,
+    alive_mask,
+    use_temporal_encoding,
+    fire_rate,
+    WEIGHTS_PATH,
+)
 
 T = transforms.Compose(
     [
         v2.ToImage(),
         v2.ToDtype(torch.float, scale=True),
         v2.ConvertImageDtype(dtype=torch.float32),
-        transforms.RandomHorizontalFlip(),
-        transforms.RandomVerticalFlip(),
+        transforms.Normalize((0.5,), (0.225,)),
     ]
 )
 
@@ -46,22 +50,26 @@ def eval_selfclass_pathmnist(
         dataset_test, shuffle=False, batch_size=32
     )
 
-    num_classes = 9
+    num_classes = len(INFO["pathmnist"]["label"])
     nca = ClassificationNCAModel(
         device,
         num_image_channels=3,
         num_hidden_channels=hidden_channels,
         hidden_size=128,
         num_classes=num_classes,
-        use_alive_mask=False,
-        fire_rate=0.5,
+        use_alive_mask=alive_mask,
+        fire_rate=fire_rate,
+        pad_noise=pad_noise,
+        use_temporal_encoding=use_temporal_encoding,
     )
     nca.load_state_dict(
-        torch.load(WEIGHTS_PATH / "selfclass_pathmnist.best.pth", weights_only=True)
+        torch.load(
+            WEIGHTS_PATH / "classification_pathmnist" / "best_model.pth",
+            weights_only=True,
+        )
     )
 
-    model_parameters = filter(lambda p: p.requires_grad, nca.parameters())
-    params = sum([np.prod(p.size()) for p in model_parameters])
+    params = nca.num_trainable_parameters()
     print(f"Trainable parameters: {params}")
     print(f"That is {4 * params / 1000} kB")
 
@@ -70,24 +78,24 @@ def eval_selfclass_pathmnist(
 
     macro_acc = torchmetrics.classification.MulticlassAccuracy(
         average="macro", num_classes=num_classes
-    )
+    ).to(device)
     micro_acc = torchmetrics.classification.MulticlassAccuracy(
         average="micro", num_classes=num_classes
-    )
+    ).to(device)
     macro_auc = torchmetrics.classification.MulticlassAUROC(
         average="macro",
         num_classes=num_classes,
-    )
+    ).to(device)
     micro_f1 = torchmetrics.classification.MulticlassF1Score(
         average="micro", num_classes=num_classes
-    )
+    ).to(device)
     for sample in tqdm(iter(loader_test)):
         x, y = sample
-        x = pad_input(x, nca, noise=True)
         x = x.float().to(device)
         steps = 72
         y_prob = nca.classify(x, steps, reduce=False)
         y = y.squeeze().to(device)
+
         macro_acc.update(y_prob, y)
         micro_acc.update(y_prob, y)
         macro_auc.update(y_prob, y)

tasks/class_medmnist/eval_class_pathmnist.py

@@ -27,6 +27,12 @@
 WEIGHTS_PATH = TASK_PATH / "weights"
 WEIGHTS_PATH.mkdir(exist_ok=True)
 
+gradient_clipping = False
+pad_noise = True
+alive_mask = False
+use_temporal_encoding = True
+fire_rate = 0.8
+
 
 def train_class_dermamnist(
     batch_size: int,
@@ -36,11 +42,6 @@ def train_class_dermamnist(
 ):
     print_NCALab_banner()
 
-    gradient_clipping = False
-    pad_noise = True
-    alive_mask = False
-    use_temporal_encoding = True
-
     comment = "DermaMNIST"
     comment += f"_hidden_{hidden_channels}"
     comment += f"_gc_{gradient_clipping}"
@@ -106,7 +107,7 @@ def train_class_dermamnist(
         num_hidden_channels=hidden_channels,
         num_classes=7,
         use_alive_mask=alive_mask,
-        fire_rate=0.8,
+        fire_rate=fire_rate,
         pad_noise=pad_noise,
         use_temporal_encoding=use_temporal_encoding,
     )

tasks/class_medmnist/train_class_dermamnist.py

@@ -14,7 +14,7 @@
 
 import click
 
-from medmnist import PathMNIST  # type: ignore[import-untyped]
+from medmnist import INFO, PathMNIST  # type: ignore[import-untyped]
 
 import torch  # type: ignore[import-untyped]
 from torchvision import transforms  # type: ignore[import-untyped]
@@ -25,6 +25,12 @@
 WEIGHTS_PATH = TASK_PATH / "weights"
 WEIGHTS_PATH.mkdir(exist_ok=True)
 
+gradient_clipping = False
+pad_noise = True
+alive_mask = False
+use_temporal_encoding = True
+fire_rate = 0.8
+
 
 def train_class_pathmnist(
     batch_size: int,
@@ -33,10 +39,6 @@ def train_class_pathmnist(
     gpu_index: int,
     lambda_activity: float,
 ):
-    gradient_clipping = False
-    pad_noise = True
-    alive_mask = False
-
     writer = SummaryWriter(
         comment=f"L.act_{lambda_activity}_c.hidden_{hidden_channels}_gc_{gradient_clipping}_noise_{pad_noise}_AM_{alive_mask}"
     )
@@ -50,6 +52,7 @@ def train_class_pathmnist(
             v2.ConvertImageDtype(dtype=torch.float32),
             transforms.RandomHorizontalFlip(),
             transforms.RandomVerticalFlip(),
+            transforms.Normalize((0.5,), (0.225,)),
         ]
     )
 
@@ -63,20 +66,21 @@ def train_class_pathmnist(
         dataset_val, shuffle=True, batch_size=32, drop_last=True
     )
 
+    num_classes = len(INFO["pathmnist"]["label"])
     nca = ClassificationNCAModel(
         device,
         num_image_channels=3,
         num_hidden_channels=hidden_channels,
-        num_classes=9,
+        num_classes=num_classes,
         use_alive_mask=alive_mask,
-        fire_rate=0.5,
+        fire_rate=fire_rate,
         lambda_activity=lambda_activity,
         filter_padding="circular",
         pad_noise=pad_noise,
     )
     trainer = BasicNCATrainer(
         nca,
-        WEIGHTS_PATH / "selfclass_pathmnist",
+        WEIGHTS_PATH / "classification_pathmnist",
         batch_repeat=2,
         max_epochs=100,
         gradient_clipping=gradient_clipping,