ncalab.models

Submodules

Classes

`BasicNCAModel`	Abstract base class for NCA models.
`BasicNCAHead`	Base class for all neural network modules.
`BasicNCAPerception`
`BasicNCARule`	NCA rule module based on a two-layer Multi-Layer-Perceptron (MLP).
`CascadeNCA`	Chain multiple instances of the same NCA model, operating at different
`ClassificationNCAModel`	Abstract base class for NCA models.
`ClassificationNCAHead`	Base class for all neural network modules.
`DepthNCAModel`	NCA model for monocular depth estimation.
`GrowingNCAModel`	NCA Model class for "growing" tasks, in which a structure is grown from a single seed pixel.
`SegmentationNCAModel`	Model used for image segmentation.

Package Contents

class ncalab.models.BasicNCAModel(device: torch.device, num_image_channels: int, num_hidden_channels: int, num_output_channels: int, plot_function: ncalab.visualization.Visual | None = None, validation_metric: str | None = None, fire_rate: float = 0.5, hidden_size: int = 128, use_alive_mask: bool = False, immutable_image_channels: bool = True, num_learned_filters: int = 2, filter_padding: str = 'reflect', use_laplace: bool = False, kernel_size: int = 3, pad_noise: bool = False, use_temporal_encoding: bool = False, rule_type: type[ncalab.models.basicNCA.basicNCArule.BasicNCARule] = BasicNCARule, training_timesteps: int | Tuple[int, int] = 100, inference_timesteps: int | Tuple[int, int] = 100)

Bases: torch.nn.Module

Abstract base class for NCA models.

BasicNCAModel is a composition of an NCA backbone model (called “rule”), and an (optional) head module for downstream tasks.

Parameters:

device – Pytorch device descriptor.
num_image_channels – Number of channels reserved for input image.
num_hidden_channels – Number of hidden channels (communication channels).
num_output_channels – Number of output channels.
fire_rate – Fire rate for stochastic weight update. Defaults to 0.5.
hidden_size – Number of neurons in hidden layer. Defaults to 128.
use_alive_mask – Whether to use alive masking (channel 3) during training. Defaults to False.
immutable_image_channels – If image channels should be fixed during inference, which is the case for most segmentation or classification problems. Defaults to True.
num_learned_filters – Number of learned filters. If zero, use two sobel filters instead. Defaults to 2.
filter_padding – Padding type to use. Might affect reliance on spatial cues. Defaults to “circular”.
use_laplace – Whether to use Laplace filter (only if num_learned_filters == 0)
kernel_size – Filter kernel size (only for learned filters)
pad_noise – Whether to pad input image tensor with noise in hidden / output channels

device

num_image_channels

num_hidden_channels

num_output_channels

num_channels

fire_rate = 0.5

hidden_size = 128

use_alive_mask = False

immutable_image_channels = True

num_learned_filters = 2

use_laplace = False

kernel_size = 3

filter_padding = 'reflect'

pad_noise = False

use_temporal_encoding = False

plot_function = None

validation_metric = None

training_timesteps = 100

inference_timesteps = 100

perception

input_vector_size

rule_type

rule

head: ncalab.models.basicNCA.basicNCAhead.BasicNCAHead | None = None

_define_rule() → ncalab.models.basicNCA.basicNCArule.BasicNCARule

prepare_input(x: torch.Tensor) → torch.Tensor

Preprocess input. Intended to be overwritten by subclass, if preprocessing is necessary.

Parameters:: [torch.Tensor] (x) – Input tensor to preprocess.
Returns:: Processed tensor.

_alive(x)

_update(x: torch.Tensor, step: int) → torch.Tensor

Compute residual cell update.

Parameters:

[torch.Tensor] (x) – Input tensor, BCWH
[int] (step) – Current timestep, required for computing temporal encoding.

Returns:

Residual cell update, BCWH.

_forward_step(x: torch.Tensor, step: int)

forward(x: torch.Tensor, steps: int = 1) → ncalab.prediction.Prediction

Parameters:

[torch.Tensor] (x) – Input image, padded along the channel dimension, BCWH.
[int] (steps) – Time steps in forward pass.

Returns [Prediction]:

Prediction object.

loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, torch.Tensor]

Compute loss. Needs to be overloaded by any subclass. Please note that the returned dict needs to hold “total” key in which the total loss is stored, which is typically a weighted sum of other losses. The total loss is backpropagated, whereas the other losses are sent to tensorboard.

Parameters:

[torch.Tensor] (label) – Input image, BCWH.
[torch.Tensor] – Ground truth, BCWH.

Returns:

Dictionary of identifiers mapped to computed losses.

finetune(freeze_head: bool = False)

Prepare model for fine tuning by freezing everything except the final layer, and setting to “train” mode.

Param:: freeze_head

metrics(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, float]

Return dict of standard evaluation metrics, given a prediction and corresponding ground truth label.

Parameters:

pred (Prediction) – Prediction.
label (torch.Tensor) – Ground truth label.

Returns:

Dict of metrics, mapped by their names.

Return type:

Dict[str, float]

predict(image: torch.Tensor, steps: int = 100) → ncalab.prediction.Prediction

Make an NCA prediction, performing multiple forward passes to yield a final result.

Parameters:

image (torch.Tensor) – Input image, BCWH.
steps (int) – Time steps

Returns:

Prediction object.

Return type:

Prediction

record(image: torch.Tensor, steps: int | None = None) → List[ncalab.prediction.Prediction]

Record predictions for all time steps and return the resulting sequence of predictions.

Parameters:: image (torch.Tensor) – Input image, BCWH.
Returns:: List of Prediction objects.
Return type:: List[Prediction]

validate(image: torch.Tensor, label: torch.Tensor, steps: int | None = None) → Tuple[Dict[str, float], ncalab.prediction.Prediction] | None

Make a prediction on an image of the validation set and return metrics computed with respect to a labelled validation image.

Parameters:

[torch.Tensor] (label) – Input image, BCWH
[torch.Tensor] – Ground truth label
[int] (steps) – Inference steps

Returns [Tuple[float, Prediction]]:

Validation metric, predicted image BCWH

_to_dict() → Dict[str, Any]

to_dict() → Dict[str, Any]

num_trainable_parameters() → int

Returns the number of trainable model parameters.

Returns:: Number of trainable parameters.
Return type:: int

class ncalab.models.BasicNCAHead

Bases: torch.nn.Module, abc.ABC

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: torch.Tensor) → torch.Tensor

Parameters:: x (torch.Tensor) – Input tensor
Returns:: NotImplemented, subclasses are required to implement this method.

abstract freeze(freeze_last: bool = True) → None

Freeze head weights.

Parameters:: freeze_last (bool, optional) – Whether to freeze the last layer (if applicable), defaults to True
Returns:: NotImplemented, subclasses are required to implement this method.

class ncalab.models.BasicNCAPerception(nca: ncalab.models.basicNCA.BasicNCAModel)

nca

_define_filters()

Define list of perception filters, based on parameters passed in constructor.

Parameters:: [int] (num_learned_filters) – Number of learned filters in perception filter bank.

perceive(x: torch.Tensor, step: int) → torch.Tensor

freeze()

class ncalab.models.BasicNCARule(device: torch.device, input_size: int, hidden_size: int, output_size: int, nonlinearity: type[torch.nn.Module] = nn.ReLU)

Bases: torch.nn.Module

NCA rule module based on a two-layer Multi-Layer-Perceptron (MLP).

Parameters:: nn (_type_) – _description_

_summary_

Parameters:

device (torch.device) – _description_
input_size (int) – _description_
hidden_size (int) – _description_
output_size (int) – _description_
nonlinearity (type[nn.Module], optional) – _description_, defaults to nn.ReLU

nonlinearity

input_size

hidden_size

output_size

device

_build_network()

_initialize_network()

Initialize network weights of the MLP.

We assume that the default initialization of the first layer is good enough. Since the final layer is purely linear and unbiased, we initalize with 0.

forward(x: torch.Tensor) → torch.Tensor

Parameters:: x (torch.Tensor) – BCWH perception vector
Returns:: BCWH residual update
Return type:: torch.Tensor

freeze(freeze_last: bool = False)

Freeze the first layer of the NCA rule network and, optionally, the final layer.

Parameters:: freeze_last (bool, optional) – _description_, defaults to False

class ncalab.models.CascadeNCA(wrapped: ncalab.models.basicNCA.BasicNCAModel, scales: List[int], steps: List[int], single_model: bool = True)

Bases: ncalab.models.basicNCA.BasicNCAModel

Chain multiple instances of the same NCA model, operating at different image scales.

The idea is to use this model as a wrapper and drop-in replacement for an existing model. For instance, if we created a model nca = SegmentationNCA(…) and all code to interface with it, we could instead write cascade = CascadeNCA(SegmentationNCA(…), scales, steps) without the need for adjusting any of the interfacing code.

This is still highly experimental. In the future, we’ll work on a cleaner interface for this.

Parameters:

wrapped (ncalab.BasicNCAModel) – Backbone model based on BasicNCAModel.
scales (List[int]) – List of scales to operate at, e.g. [4, 2, 1].
steps (List[int]) – List of number of NCA inference time steps.
single_model (bool) – Only train a single instance of the NCA model

loss

Compute loss. Needs to be overloaded by any subclass. Please note that the returned dict needs to hold “total” key in which the total loss is stored, which is typically a weighted sum of other losses. The total loss is backpropagated, whereas the other losses are sent to tensorboard.

Parameters:

[torch.Tensor] (label) – Input image, BCWH.
[torch.Tensor] – Ground truth, BCWH.

Returns:

Dictionary of identifiers mapped to computed losses.

finetune

Prepare model for fine tuning by freezing everything except the final layer, and setting to “train” mode.

Param:: freeze_head

prepare_input

Preprocess input. Intended to be overwritten by subclass, if preprocessing is necessary.

Parameters:: [torch.Tensor] (x) – Input tensor to preprocess.
Returns:: Processed tensor.

head

wrapped

scales

steps

single_model = True

models: List[ncalab.models.basicNCA.BasicNCAModel]

forward(x: torch.Tensor, *args, **kwargs) → ncalab.prediction.Prediction

Parameters:

x (torch.Tensor) – Input image tensor, BCWH.
steps (torch.Tensor) – Unused, as steps are defined in constructor.

Returns:

Prediction object

Return type:

Prediction

record(image: torch.Tensor, steps: int | None = None) → List[ncalab.prediction.Prediction]

Records predictions for all time steps and returns the resulting sequence of predictions.

Takes care of scaling the image in between steps.

Parameters:: image (torch.Tensor) – Input image, BCWH.
Returns:: List of Prediction objects.
Return type:: List[Prediction]

validate(image: torch.Tensor, label: torch.Tensor, steps: int | None = None) → Tuple[Dict[str, float], ncalab.prediction.Prediction] | None

Validation method.

Takes care of scaling the input image and label on each scale, and calls the respective validation method of the wrapped model.

Parameters:

image – Input image tensor, BCWH.
label – Ground truth.
steps (int) – Unused, as steps are defined in constructor.

Returns:

Dict of metrics and prediction

Return type:

Optional[Tuple[Dict[str, float], Prediction]]

class ncalab.models.ClassificationNCAModel(device: torch.device, num_image_channels: int, num_hidden_channels: int, num_classes: int, fire_rate: float = 0.8, hidden_size: int = 128, use_alive_mask: bool = False, pixel_wise_loss: bool = False, num_learned_filters: int = 2, filter_padding: str = 'reflect', use_laplace: bool = False, kernel_size: int = 3, pad_noise: bool = False, use_temporal_encoding: bool = False, use_classifier: bool = True, class_names: List[str] | None = None, avg_pool_size: int = 5, **kwargs)

Bases: ncalab.models.basicNCA.BasicNCAModel

Abstract base class for NCA models.

BasicNCAModel is a composition of an NCA backbone model (called “rule”), and an (optional) head module for downstream tasks.

Parameters:

device – Pytorch device descriptor.
num_image_channels – _description_
num_hidden_channels – _description_
num_classes – _description_
fire_rate – Fire rate for stochastic weight update. Defaults to 0.8.
hidden_size – Number of neurons in hidden layer. Defaults to 128.
use_alive_mask – Whether to use alive masking (channel 3) during training. Defaults to False.
pixel_wise_loss – Whether a prediction per pixel is desired, like in self-classifying MNIST. Defaults to False.
num_learned_filters – Number of learned filters. If zero, use two sobel filters instead. Defaults to 2.
filter_padding – Padding type to use. Might affect reliance on spatial cues. Defaults to “circular”.
pad_noise – Whether to pad input image tensor with noise in hidden / output channels

_num_classes

pixel_wise_loss = False

use_classifier = True

avg_pool_size = 5

property num_classes: int

classify(image: torch.Tensor, steps: int = 100, reduce: bool = False) → torch.Tensor

Predict classification for an input image.

Parameters:

image – Input image.
steps – Inference steps. Defaults to 100.
reduce – Return a single softmax probability. Defaults to False.

Returns:

Single class index or vector of logits.

loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, torch.Tensor]

Return the classification loss.

Parameters:

pred – Prediction.
label – Ground truth.

Returns:

Dictionary of identifiers mapped to computed losses.

metrics(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, float]

Return dict of standard evaluation metrics.

Parameters:

pred (Prediction) – Prediction
label (Tensor) – Ground truth label.

class ncalab.models.ClassificationNCAHead(num_hidden_channels: int, num_classes: int, device: torch.device, avg_pool_size: int = 3, hidden_size: int = 64)

Bases: ncalab.models.basicNCA.basicNCAhead.BasicNCAHead

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

num_hidden_channels

num_classes

device

avg_pool_size = 3

classifier

forward(x)

Parameters:: x (torch.Tensor) – Input tensor
Returns:: NotImplemented, subclasses are required to implement this method.

freeze(freeze_last: bool = False)

Freeze head weights.

Parameters:: freeze_last (bool, optional) – Whether to freeze the last layer (if applicable), defaults to True
Returns:: NotImplemented, subclasses are required to implement this method.

class ncalab.models.DepthNCAModel(device: torch.device, num_image_channels: int = 3, num_hidden_channels: int = 18, fire_rate: float = 0.8, hidden_size: int = 128, num_learned_filters: int = 2, pad_noise: bool = False, **kwargs)

Bases: ncalab.models.basicNCA.BasicNCAModel

NCA model for monocular depth estimation.

Parameters:

device – Pytorch device descriptor.
num_image_channels – Number of channels reserved for input image.
num_hidden_channels – Number of hidden channels (communication channels).
num_output_channels – Number of output channels.
fire_rate – Fire rate for stochastic weight update. Defaults to 0.5.
hidden_size – Number of neurons in hidden layer. Defaults to 128.
use_alive_mask – Whether to use alive masking (channel 3) during training. Defaults to False.
immutable_image_channels – If image channels should be fixed during inference, which is the case for most segmentation or classification problems. Defaults to True.
num_learned_filters – Number of learned filters. If zero, use two sobel filters instead. Defaults to 2.
filter_padding – Padding type to use. Might affect reliance on spatial cues. Defaults to “circular”.
use_laplace – Whether to use Laplace filter (only if num_learned_filters == 0)
kernel_size – Filter kernel size (only for learned filters)
pad_noise – Whether to pad input image tensor with noise in hidden / output channels

vignette = None

loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, torch.Tensor]

Parameters:

image – Input image, BCWH.
label – Ground truth.

Returns:

Dictionary of identifiers mapped to computed losses.

metrics(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, float]

Return dict of standard evaluation metrics. Needs to include special item ‘prediction’, containing the predicted image (all channels).

Parameters:

pred
label – Ground truth label.

Returns [Dict]:

Dict of metrics, mapped by their names.

class ncalab.models.GrowingNCAModel(device: torch.device, num_image_channels: int = 4, num_hidden_channels: int = 16, fire_rate: float = 0.5, hidden_size: int = 128, use_alive_mask: bool = False, **kwargs)

Bases: ncalab.models.basicNCA.BasicNCAModel

NCA Model class for “growing” tasks, in which a structure is grown from a single seed pixel.

This specialization of the BasicNCAModel has some interesting properties. For instance, it has no output channels, as the growing task directly manipulates the input image channels.

Parameters:

[torch.device] (device) – Pytorch device descriptor.
[int] (hidden_size) – Number of channels reserved for input image. Defaults to 4.
[int] – Number of hidden channels (communication channels). Defaults to 16.
[float] (fire_rate) – Stochastic weight update. Defaults to 0.5.
[int] – Default number of nodes in hidden layer. Defaults to 128.
[bool] (use_alive_mask) – Whether to use alive masking. Defaults to False.

loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, torch.Tensor]

Implements a simple MSE loss between target and prediction.

Parameters:

pred – Prediction
label – Target

Returns [Tensor]:

MSE Loss

validate(image: torch.Tensor, label: torch.Tensor, steps: int | None = None) → Tuple[Dict[str, float], ncalab.prediction.Prediction] | None: We typically don’t validate during training of Growing NCA, because there is only a single sample in the training set.

make_seed(width: int, height: int) → torch.Tensor

grow(seed: torch.Tensor, steps: int = 100) → List[numpy.ndarray]

Run the growth process and return the resulting output sequence.

Parameters:

[torch.Tensor] (seed) – Seed image, can be generated through make_seed.
[int] (steps) – Number of inference steps. Defaults to 100.

Returns [List[np.ndarray]]:

Sequence of output images.

class ncalab.models.SegmentationNCAModel(device: torch.device, num_image_channels: int = 3, num_hidden_channels: int = 16, num_classes: int = 1, fire_rate: float = 0.8, hidden_size: int = 128, num_learned_filters: int = 2, pad_noise: bool = True, filter_padding: str = 'reflect', **kwargs)

Bases: ncalab.models.basicNCA.BasicNCAModel

Model used for image segmentation.

Uses Dice score as the default validation metric. Currently, only binary segmentation masks are supported.

Parameters:

[torch.device] (device) – Compute device.
[int] (learned_filters) – Number of image channels. Defaults to 3.
[int] – Number of hidden channels. Defaults to 16.
[int] – Number of classes. Defaults to 1.
[float] (fire_rate) – NCA fire rate. Defaults to 0.8.
[int] – Number of neurons in hidden layer. Defaults to 128.
[int] – Number of learned filters. If 0, use sobel. Defaults to 2.
[bool] (pad_noise) – Whether to pad input images with noise. Defaults to True.
[str] (filter_padding) – Padding type to use. Might affect reliance on spatial cues. Defaults to “circular”.

num_classes = 1

loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, torch.Tensor]

Compute Dice + BCE loss.

Parameters:

pred – Prediction.
label – Ground truth.

Returns:

Dictionary of identifiers mapped to computed losses.

metrics(pred: ncalab.prediction.Prediction, label: torch.Tensor) → Dict[str, float]

Return dict of standard evaluation metrics.

Parameters:

[torch.Tensor] (label) – Predicted image.
[torch.Tensor] – Ground truth label.