ncalab.uncertainty

Submodules

ncalab.uncertainty.uncertainty_estimator

Classes

`BasicNCAModel`	Abstract base class for NCA models.
`Prediction`	Stores the result of an NCA prediction, including the number of steps it took.
`UncertaintyEstimator`
`NQM`	Variance over multiple predictions.
`MCMC`	Markov-Chain Monte Carlo

Package Contents

class ncalab.uncertainty.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.uncertainty.Prediction(model, steps: int, output_image: torch.Tensor, head_prediction: torch.Tensor | None = None)

Stores the result of an NCA prediction, including the number of steps it took.

Sequences are typically stored by BasicNCAModel’s “record” function, and are returned as a list of Prediction objects.

Constructor is typically not called explicitly. Rather, the forward pass of BasicNCAModel (and its subclasses) is responsible for filling its attributes.

Parameters:

model (ncalab.BasicNCAModel) – Reference to model used for prediction.
steps (int) – Number of steps taken for the prediction.
output_image (torch.Tensor) – Output image tensor.

model

steps

output_image

_output_array: numpy.ndarray | None = None

head_prediction = None

_head_prediction_array: numpy.ndarray | None = None

property image_channels: torch.Tensor

Convenience property to access the image channels as a Tensor.

Returns:: BCWH Tensor
Return type:: torch.Tensor

property hidden_channels: torch.Tensor

Convenience property to access the hidden channels as a Tensor.

Returns:: BCWH Tensor
Return type:: torch.Tensor

property output_channels: torch.Tensor

Convenience property to access the output channels as a Tensor.

Returns:: BCWH Tensor
Return type:: torch.Tensor

property output_array: numpy.ndarray

Convenience property to access the whole output image in the format of a numpy array. Brings the entire tensor to CPU on demand, and only at the first call.

Returns:: Numpy array in BCWH format
Return type:: np.ndarray

property image_channels_np: numpy.ndarray

Convenience property to access the output image channels in the format of a numpy array. Brings the entire tensor to CPU on demand, and only at the first call.

Returns:: Numpy array in BCWH format
Return type:: np.ndarray

property hidden_channels_np: numpy.ndarray

Convenience property to access the hidden image channels in the format of a numpy array. Brings the entire tensor to CPU on demand, and only at the first call.

Returns:: Numpy array in BCWH format
Return type:: np.ndarray

property output_channels_np: numpy.ndarray

Convenience property to access the image’s output channels in the format of a numpy array. Brings the entire tensor to CPU on demand, and only at the first call.

Returns:: Numpy array in BCWH format
Return type:: np.ndarray

property head_prediction_array: numpy.ndarray | None

class ncalab.uncertainty.UncertaintyEstimator(nca: ncalab.models.BasicNCAModel)

Parameters:: nca (BasicNCAModel) – Trained NCA model

nca

_estimate(image: torch.Tensor) → Tuple[torch.Tensor, List[ncalab.prediction.Prediction]]

Internal uncertainty estimation method.

Parameters:: image (torch.Tensor) – Input image sample
Returns:: Float tensor of uncertainty heatmap (BCWH), predictions, reduced uncertainty score
Return type:: Tuple[torch.Tensor, List[Prediction], float]

estimate(image: torch.Tensor, reduce: str = 'mean') → Tuple[torch.Tensor, List[ncalab.prediction.Prediction], torch.Tensor]

Estimate predictive uncertainty.

Parameters:

image (torch.Tensor) – Input image sample
reduce (str) – Reduction strategy, defaults to “mean”.

Returns:

Float tensor of uncertainty heatmap (BCWH), final prediction, reduced uncertainty score for batch (BC)

Return type:

Tuple[torch.Tensor, List[Prediction], torch.Tensor]

__call__(*args, **kwargs)

class ncalab.uncertainty.NQM(nca: ncalab.models.BasicNCAModel, N: int = 10, normalize=False)

Bases: UncertaintyEstimator

Variance over multiple predictions.

Parameters:: nca (BasicNCAModel) – Trained NCA model

N = 10

normalize = False

_estimate(image: torch.Tensor) → Tuple[torch.Tensor, List[ncalab.prediction.Prediction]]

Internal uncertainty estimation method.

Parameters:: image (torch.Tensor) – Input image sample
Returns:: Float tensor of uncertainty heatmap (BCWH), predictions, reduced uncertainty score
Return type:: Tuple[torch.Tensor, List[Prediction], float]

class ncalab.uncertainty.MCMC(nca: ncalab.models.BasicNCAModel, N_last: int = 10, normalize=False)

Bases: UncertaintyEstimator

Markov-Chain Monte Carlo

Parameters:: nca (BasicNCAModel) – Trained NCA model

N_last = 10

normalize = False

_estimate(image: torch.Tensor) → Tuple[torch.Tensor, List[ncalab.prediction.Prediction]]

Internal uncertainty estimation method.

Parameters:: image (torch.Tensor) – Input image sample
Returns:: Float tensor of uncertainty heatmap (BCWH), predictions, reduced uncertainty score
Return type:: Tuple[torch.Tensor, List[Prediction], float]