let forward() method return Prediction object

Author: henrykrumb

ncalab/__init__.py

@@ -1,8 +1,10 @@
+from .autostepper import AutoStepper  # noqa: F401,F403
 from .data import GrowingNCADataset  # noqa: F401
 from .experiment import *  # noqa: F401,F403
 from .losses import *  # noqa: F401,F403
 from .models import *  # noqa: F401,F403
 from .paths import *  # noqa: F401,F403
+from .prediction import *  # noqa: F401,F403
 from .search import *  # noqa: F401,F403
 from .training import *  # noqa: F401,F403
 from .utils import *  # noqa: F401,F403

ncalab/autostepper.py

@@ -1,5 +1,7 @@
 import logging
 
+import torch
+
 
 class AutoStepper:
     """
@@ -20,8 +22,8 @@ def __init__(
         :param min_steps [int]: Minimum number of timesteps to always execute. Defaults to 10.
         :param max_steps [int]: Terminate after maximum number of steps. Defaults to 100.
         :param plateau [int]: _description_. Defaults to 5.
-        :param verbose [bool]: Whether to communicate. Defaults to False.
-         threshold (float, optional): _description_. Defaults to 1e-2.
+        :param verbose [bool]: Whether to log interruption to stdout. Defaults to False.
+        :param threshold [float]: Score threshold. Defaults to 1e-2.
         """
         assert min_steps >= 1
         assert plateau >= 1
@@ -32,22 +34,38 @@ def __init__(
         self.verbose = verbose
         self.threshold = threshold
         self.cooldown = 0
+        # invariant: auto_min_steps > 0, so both of these will be defined when used
+        self.hidden_i: torch.Tensor | None = None
+        self.hidden_i_1: torch.Tensor | None = None
+
+    def score(self) -> torch.Tensor:
+        assert self.hidden_i is not None
+        assert self.hidden_i_1 is not None
+        # normalized absolute difference between two hidden states
+        return (self.hidden_i - self.hidden_i_1).abs().sum() / torch.numel(
+            self.hidden_i
+        )
 
-    def check(self, step, score):
+    def check(self, step: int) -> bool:
         """
-        _summary_
+        Checks whether to interrupt inference after the current step.
 
-        :param score: _description_
-        :type score: _type_
-        :return: _description_
-        :rtype: _type_
+        :param score [int]: Current NCA inference step.
+        :return [bool]: Whether to interrupt inference after the current step.
         """
-        if score >= self.threshold:
-            self.cooldown = 0
-        else:
-            self.cooldown += 1
-        if self.cooldown >= self.plateau:
-            if self.verbose:
-                logging.info(f"Breaking after {step} steps.")
-            return True
-        return False
+        with torch.no_grad():
+            if step < self.min_steps:
+                return False
+            if step >= self.max_steps:
+                return True
+            if self.hidden_i is None or self.hidden_i_1 is None:
+                return False
+            if self.score() >= self.threshold:
+                self.cooldown = 0
+            else:
+                self.cooldown += 1
+            if self.cooldown >= self.plateau:
+                if self.verbose:
+                    logging.info(f"Breaking after {step} steps.")
+                return True
+            return False

ncalab/export/header.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
-import torch
+import os
 from pathlib import PosixPath, Path
 
+import torch
+
 from ..models.basicNCA import BasicNCAModel
 
 
@@ -34,12 +36,12 @@ def export_header(
 ):
     ## prepare preamble
     # guard
-    preamble = "#pragma once\n"
+    preamble = f"#pragma once{os.linesep}"
     # add imports if any
     if imports:
         for header in imports:
-            preamble += f'#include "{header}"'
-    preamble += "\n\n"
+            preamble += f'#include "{header}"{os.linesep}'
+    preamble += f"{os.linesep}{os.linesep}"
 
     with open(outfile, "w") as f:
         f.write(preamble)

ncalab/hooks/__init__.py

@@ -0,0 +1 @@
+from .hook import Hook  # noqa: F401

ncalab/hooks/hook.py

@@ -0,0 +1,18 @@
+import torch
+
+
+class Hook:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def pre_forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def post_forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def pre_perceive(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def pre_update(self, x: torch.Tensor) -> torch.Tensor:
+        return x

ncalab/hooks/hook_output_noise.py

@@ -0,0 +1,19 @@
+import torch
+
+from .hook import Hook
+
+class OutputNoiseHook(Hook):
+    def __init__(self, dx_noise: float):
+        self.dx_noise = dx_noise
+
+    def pre_forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def post_forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def pre_perceive(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def pre_update(self, x: torch.Tensor) -> torch.Tensor:
+        return x

ncalab/models/basicNCA.py

@@ -8,6 +8,7 @@
 import torch.nn.functional as F  # type: ignore[import-untyped]
 
 from ..autostepper import AutoStepper
+from ..prediction import Prediction
 from ..utils import pad_input
 
 
@@ -22,6 +23,8 @@ def __init__(
         num_image_channels: int,
         num_hidden_channels: int,
         num_output_channels: int,
+        plot_function: Optional[Callable] = None,
+        validation_metric: Optional[str] = None,
         fire_rate: float = 0.5,
         hidden_size: int = 128,
         use_alive_mask: bool = False,
@@ -74,14 +77,15 @@ def __init__(
         self.pad_noise = pad_noise
         self.autostepper = autostepper
         self.use_temporal_encoding = use_temporal_encoding
-
-        # set by subclassing functions
-        self.plot_function: Optional[Callable] = None
-        self.validation_metric: Optional[str] = None
+        self.plot_function = plot_function
+        self.validation_metric = validation_metric
 
         self._define_filters(num_learned_filters)
 
         # define model structure
+        self._define_network()
+
+    def _define_network(self):
         input_vector_size = self.num_channels * (self.num_filters + 1)
         if self.use_temporal_encoding:
             input_vector_size += 1
@@ -103,7 +107,7 @@ def __init__(
                 padding=0,
                 kernel_size=1,
             ),
-        ).to(device)
+        ).to(self.device)
 
         # initialize final layer with 0
         with torch.no_grad():
@@ -137,8 +141,7 @@ def _define_filters(self, num_learned_filters: int):
             sobel_x = np.outer([1, 2, 1], [-1, 0, 1]) / 8.0
             sobel_y = sobel_x.T
             laplace = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
-            self.filters.append(sobel_x)
-            self.filters.append(sobel_y)
+            self.filters.extend([sobel_x, sobel_y])
             if self.use_laplace:
                 self.filters.append(laplace)
             self.num_filters = len(self.filters)
@@ -154,7 +157,7 @@ def prepare_input(self, x: torch.Tensor) -> torch.Tensor:
         """
         return x
 
-    def __alive(self, x):
+    def _alive(self, x):
         mask = (
             F.max_pool2d(
                 x[:, 3, :, :],
@@ -190,11 +193,12 @@ def _perceive_with(x, weight):
         dx = torch.cat(perception, 1)
         return dx
 
-    def _update(self, x: torch.Tensor, step):
+    def _update(self, x: torch.Tensor, step: int) -> torch.Tensor:
         """
         Compute residual cell update.
 
         :param x [torch.Tensor]: Input tensor, BCWH
+        :param step [int]: Current timestep, required for computing temporal encoding.
         """
         assert x.shape[1] == self.num_channels
 
@@ -218,40 +222,33 @@ def forward(
         self,
         x: torch.Tensor,
         steps: int = 1,
-    ) -> torch.Tensor | Tuple[torch.Tensor, int]:
+    ) -> Prediction:
         """
         :param x [torch.Tensor]: Input image, padded along the channel dimension, BCWH.
         :param steps [int]: Time steps in forward pass.
 
-        :returns: Output image (BCWH)
+        :returns [Prediction]: Prediction object.
         """
         if self.autostepper is None:
             for step in range(steps):
                 dx = self._update(x, step)
                 x = x + dx
-            return x, steps
 
-        # invariant: auto_min_steps > 0, so both of these will be defined when used
-        hidden_i: torch.Tensor | None = None
-        hidden_i_1: torch.Tensor | None = None
+                # Alive masking
+                if self.use_alive_mask:
+                    life_mask = self._alive(x)
+                    life_mask = life_mask
+                    x = x.permute(1, 0, 2, 3)  # B C W H --> C B W H
+                    x = x * life_mask.float()
+                    x = x.permute(1, 0, 2, 3)  # C B W H --> B C W H
+            return Prediction(self, steps, x)
+
+
         for step in range(self.autostepper.max_steps):
-            with torch.no_grad():
-                if (
-                    step >= self.autostepper.min_steps
-                    and hidden_i is not None
-                    and hidden_i_1 is not None
-                ):
-                    # normalized absolute difference between two hidden states
-                    score = (hidden_i - hidden_i_1).abs().sum() / (
-                        hidden_i.shape[0]
-                        * hidden_i.shape[1]
-                        * hidden_i.shape[2]
-                        * hidden_i.shape[3]
-                    )
-                    if self.autostepper.check(step, score):
-                        return x, step
+            if self.autostepper.check(step):
+                return Prediction(self, step, x)
             # save previous hidden state
-            hidden_i_1 = x[
+            self.autostepper.hidden_i_1 = x[
                 :,
                 self.num_image_channels : self.num_image_channels
                 + self.num_hidden_channels,
@@ -264,20 +261,21 @@ def forward(
 
             # Alive masking
             if self.use_alive_mask:
-                life_mask = self.__alive(x)
+                life_mask = self._alive(x)
                 life_mask = life_mask
                 x = x.permute(1, 0, 2, 3)  # B C W H --> C B W H
                 x = x * life_mask.float()
                 x = x.permute(1, 0, 2, 3)  # C B W H --> B C W H
+
             # set current hidden state
-            hidden_i = x[
+            self.autostepper.hidden_i = x[
                 :,
                 self.num_image_channels : self.num_image_channels
                 + self.num_hidden_channels,
                 :,
                 :,
             ]
-        return x, self.autostepper.max_steps
+        return Prediction(self, self.autostepper.max_steps, x)
 
     def loss(self, image: torch.Tensor, label: torch.Tensor) -> Dict[str, torch.Tensor]:
         """
@@ -317,11 +315,11 @@ def metrics(self, pred: torch.Tensor, label: torch.Tensor) -> Dict[str, float]:
         """
         return {}
 
-    def predict(self, image: torch.Tensor, steps: int = 100) -> torch.Tensor:
+    def predict(self, image: torch.Tensor, steps: int = 100) -> Prediction:
         """
         :param image [torch.Tensor]: Input image, BCWH.
 
-        :returns [torch.Tensor]: Output image, BCWH
+        :returns [Prediction]: Prediction object.
         """
         assert steps >= 1
         assert image.shape[1] <= self.num_channels
@@ -330,19 +328,22 @@ def predict(self, image: torch.Tensor, steps: int = 100) -> torch.Tensor:
             x = image.clone()
             x = pad_input(x, self, noise=self.pad_noise)
             x = self.prepare_input(x)
-            x, _ = self.forward(x, steps=steps)  # type: ignore[assignment]
-            return x
+            prediction = self.forward(x, steps=steps)
+            return prediction
 
     def validate(
         self, image: torch.Tensor, label: torch.Tensor, steps: int
-    ) -> Optional[Tuple[Dict[str, float], torch.Tensor]]:
+    ) -> Optional[Tuple[Dict[str, float], Prediction]]:
         """
+        Make a prediction on an image of the validation set and return metrics computed
+        with respect to a labelled validation image.
+
         :param image [torch.Tensor]: Input image, BCWH
         :param label [torch.Tensor]: Ground truth label
         :param steps [int]: Inference steps
 
-        :returns [Tuple[float, torch.Tensor]]: Validation metric, predicted image BCWH
+        :returns [Tuple[float, Prediction]]: Validation metric, predicted image BCWH
         """
-        pred = self.predict(image.to(self.device), steps=steps)
-        metrics = self.metrics(pred, label.to(self.device))
-        return metrics, pred
+        prediction = self.predict(image.to(self.device), steps=steps)
+        metrics = self.metrics(prediction.output_image, label.to(self.device))
+        return metrics, prediction