ncalab.models

Submodules

Classes

`AbstractNCAModel`	Abstract base class for NCA models.
`AbstractNCAHead`	Base class for all neural network modules.
`BasicNCAPerception`
`AbstractNCARule`	Base class for all neural network modules.
`MLPNCARule`	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.AbstractNCAModel(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 = 0, filter_padding: Literal['zero', 'reflect', 'replicate', 'circular'] = 'reflect', use_laplace: bool = False, kernel_size: int = 3, pad_noise: bool = False, use_temporal_encoding: bool = False, rule_type: type[ncalab.models.basicNCA.abstractNCArule.AbstractNCARule] = MLPNCARule, rule_args=None, training_timesteps: int | Tuple[int, int] = 100, inference_timesteps: int | Tuple[int, int] = 100)

Bases: torch.nn.Module, abc.ABC

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.
validation_metric
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
use_temporal_encoding
rule_type
rule_args
training_timesteps
inference_timesteps

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 = 0

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_args = None

rule

head: ncalab.models.basicNCA.abstractNCAhead.AbstractNCAHead | None = None

metrics: Dict[str, torchmetrics.Metric]

_define_rule() → ncalab.models.basicNCA.abstractNCArule.AbstractNCARule

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.

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

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

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

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

Parameters:

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

Returns:

Prediction object.

Return type:

Prediction

record(image: torch.Tensor, steps: int | Tuple[int, 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(dataloader: torch.utils.data.DataLoader, steps: int | None = None) → Tuple[Dict[str, float], List[ncalab.prediction.Prediction]]

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

Parameters:

[torch.utils.data.DataLoader] (dataloader) – Dataloader for validation images
[int] (steps) – Inference steps

Returns [Tuple[float, List[Prediction]]]:

Validation metric, predicted image BCWH

_to_dict() → Dict[str, Any]

to_dict() → Dict[str, Any]

classmethod from_dict(d: Dict[str, Any])

num_trainable_parameters() → int

Returns the number of trainable model parameters.

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

save(path: str | os.PathLike)

static load(model: AbstractNCAModel, path: str | os.PathLike) → AbstractNCAModel

post_prediction(prediction: ncalab.prediction.Prediction) → ncalab.prediction.Prediction

class ncalab.models.AbstractNCAHead

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.

optimizer = None

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

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

abstractmethod 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.AbstractNCAModel)

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.AbstractNCARule(device: torch.device, input_size: int, hidden_size: int, output_size: int)

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.

device

input_size

hidden_size

output_size

abstractmethod freeze(freeze_last: bool = False) → None

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

Bases: ncalab.models.basicNCA.abstractNCArule.AbstractNCARule

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

Parameters:

nn (_type_) – _description_
device (torch.device) – Compute device
input_size (int) – Input neurons
hidden_size (int) – Hidden neurons
output_size (int) – Output neurons
nonlinearity (type[nn.Module], optional) – Activation function, defaults to nn.ReLU

nonlinearity

_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.AbstractNCAModel, scales: List[int], steps: List[int], single_model: bool = True)

Bases: ncalab.models.basicNCA.AbstractNCAModel

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.AbstractNCAModel) – Backbone model based on AbstractNCAModel.
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

metrics

wrapped

scales

steps

single_model = True

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

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 | Tuple[int, 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]

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: Literal['zero', 'reflect', 'replicate', 'circular'] = '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 = 8, lambda_hidden: float = 0, **kwargs)

Bases: ncalab.models.basicNCA.AbstractNCAModel

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 = 8

lambda_hidden = 0

metrics

focal_loss

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.

post_prediction(prediction: ncalab.prediction.Prediction) → ncalab.prediction.Prediction

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

Bases: ncalab.models.basicNCA.abstractNCAhead.AbstractNCAHead

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

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.AbstractNCAModel

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.
validation_metric
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
use_temporal_encoding
rule_type
rule_args
training_timesteps
inference_timesteps

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.

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, lambda_hidden: float = 0.0, **kwargs)

Bases: ncalab.models.basicNCA.AbstractNCAModel

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.

lambda_hidden = 0.0

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

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 = False, filter_padding: Literal['zero', 'reflect', 'replicate', 'circular'] = 'circular', lambda_hidden: float = 0.001, **kwargs)

Bases: ncalab.models.basicNCA.AbstractNCAModel

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

metrics

lambda_hidden = 0.001

bce_loss

dice_loss

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

Compute Dice loss.

Parameters:

pred – Prediction.
label – Ground truth.

Returns:

Dictionary of identifiers mapped to computed losses.

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

post_prediction(prediction: ncalab.prediction.Prediction) → ncalab.prediction.Prediction