create evaluation script for selfclass MNIST

Author: henrykrumb

tasks/selfclass_mnist/eval_selfclass_mnist.py

@@ -5,10 +5,9 @@
 
 from ncalab import (
     ClassificationNCAModel,
-    BasicNCATrainer,
-    WEIGHTS_PATH,
-    show_batch_binary_image_classification,
     get_compute_device,
+    pad_input,
+    WEIGHTS_PATH
 )
 
 import click
@@ -21,8 +20,37 @@
 from torchvision import transforms  # type: ignore[import-untyped]
 
 
+def print_MNIST_digit(image, prediction):
+    BG = {
+        0: "black",
+        1: "red",
+        2: "cyan",
+        3: "green",
+        4: "magenta",
+        5: "yellow",
+        6: "green",
+        7: "white",
+        8: "blue",
+        9: "red"
+    }
+    FG = {
+        0: "white",
+        6: "red",
+        7: "black",
+    }
+    for y in range(14):
+        for x in range(14):
+            if image[y * 2, x * 2] < 0.3:
+                click.secho("  ", nl=False, fg="black", bg="black")
+                continue
+            n = prediction[y * 2, x * 2].detach().cpu()
+            n = int(torch.argmax(n))
+            click.secho(f" {n}", nl=False, fg=FG.get(n, "black"), bg=BG.get(n, "white"))
+        click.secho()
+
+
 def eval_selfclass_mnist(
-    batch_size: int, hidden_channels: int, gpu: bool, gpu_index: int
+    hidden_channels: int, gpu: bool, gpu_index: int
 ):
     mnist_test = MNIST(
         "mnist",
@@ -32,7 +60,7 @@ def eval_selfclass_mnist(
     )
 
     loader_test = torch.utils.data.DataLoader(
-        mnist_test, shuffle=True, batch_size=batch_size
+        mnist_test, shuffle=True, batch_size=1
     )
 
     device = get_compute_device(f"cuda:{gpu_index}" if gpu else "cpu")
@@ -44,6 +72,24 @@ def eval_selfclass_mnist(
         num_classes=10,
         pixel_wise_loss=True,
     )
+    nca.load_state_dict(torch.load(WEIGHTS_PATH / "selfclass_mnist.pth", weights_only=True))
+    nca.eval()
+
+    i = 1
+    for image, _ in loader_test:
+        if i == 0:
+            break
+        x = image.clone()
+        x = pad_input(x, nca, noise=False)
+        x = x.permute(0, 2, 3, 1).to(device)
+
+        prediction = nca(x, steps=50)[0]
+        prediction = prediction[..., nca.num_image_channels + nca.num_hidden_channels :]
+        print_MNIST_digit(image[0, 0], prediction)
+
+        if i != 1:
+            click.secho("-" * 28 * 2)
+        i -= 1
 
 
 @click.command()
@@ -54,7 +100,7 @@ def eval_selfclass_mnist(
 @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, gpu_index):
+def main(hidden_channels, gpu, gpu_index):
     eval_selfclass_mnist(
         hidden_channels=hidden_channels,
         gpu=gpu,

tasks/selfclass_mnist/train_selfclass_mnist.py

@@ -84,6 +84,7 @@ def train_selfclass_mnist(
     )
     trainer.train_basic_nca(
         loader_train,
+        # Validation is broken here. We're working on it!
         # loader_val,
         summary_writer=writer,
         plot_function=show_batch_binary_image_classification,