ncalab.models ============= .. py:module:: ncalab.models Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/ncalab/models/applications/index /autoapi/ncalab/models/basicNCA/index /autoapi/ncalab/models/wrappers/index Classes ------- .. autoapisummary:: ncalab.models.AbstractNCAModel ncalab.models.AbstractNCAHead ncalab.models.BasicNCAPerception ncalab.models.AbstractNCARule ncalab.models.MLPNCARule ncalab.models.CascadeNCA ncalab.models.ClassificationNCAModel ncalab.models.ClassificationNCAHead ncalab.models.DepthNCAModel ncalab.models.GrowingNCAModel ncalab.models.SegmentationNCAModel Package Contents ---------------- .. py:class:: AbstractNCAModel(device: torch.device, num_image_channels: int, num_hidden_channels: int, num_output_channels: int, plot_function: Optional[ncalab.visualization.Visual] = None, validation_metric: Optional[str] = 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: :py:obj:`torch.nn.Module`, :py:obj:`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. :param device: Pytorch device descriptor. :param num_image_channels: Number of channels reserved for input image. :param num_hidden_channels: Number of hidden channels (communication channels). :param num_output_channels: Number of output channels. :param validation_metric: :param fire_rate: Fire rate for stochastic weight update. Defaults to 0.5. :param hidden_size: Number of neurons in hidden layer. Defaults to 128. :param use_alive_mask: Whether to use alive masking (channel 3) during training. Defaults to False. :param immutable_image_channels: If image channels should be fixed during inference, which is the case for most segmentation or classification problems. Defaults to True. :param num_learned_filters: Number of learned filters. If zero, use two sobel filters instead. Defaults to 2. :param filter_padding: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular". :param use_laplace: Whether to use Laplace filter (only if num_learned_filters == 0) :param kernel_size: Filter kernel size (only for learned filters) :param pad_noise: Whether to pad input image tensor with noise in hidden / output channels :param use_temporal_encoding: :param rule_type: :param rule_args: :param training_timesteps: :param inference_timesteps: .. py:attribute:: device .. py:attribute:: num_image_channels .. py:attribute:: num_hidden_channels .. py:attribute:: num_output_channels .. py:attribute:: num_channels .. py:attribute:: fire_rate :value: 0.5 .. py:attribute:: hidden_size :value: 128 .. py:attribute:: use_alive_mask :value: False .. py:attribute:: immutable_image_channels :value: True .. py:attribute:: num_learned_filters :value: 0 .. py:attribute:: use_laplace :value: False .. py:attribute:: kernel_size :value: 3 .. py:attribute:: filter_padding :value: 'reflect' .. py:attribute:: pad_noise :value: False .. py:attribute:: use_temporal_encoding :value: False .. py:attribute:: plot_function :value: None .. py:attribute:: validation_metric :value: None .. py:attribute:: training_timesteps :value: 100 .. py:attribute:: inference_timesteps :value: 100 .. py:attribute:: perception .. py:attribute:: input_vector_size .. py:attribute:: rule_type .. py:attribute:: rule_args :value: None .. py:attribute:: rule .. py:attribute:: head :type: ncalab.models.basicNCA.abstractNCAhead.AbstractNCAHead | None :value: None .. py:attribute:: metrics :type: Dict[str, torchmetrics.Metric] .. py:method:: _define_rule() -> ncalab.models.basicNCA.abstractNCArule.AbstractNCARule .. py:method:: prepare_input(x: torch.Tensor) -> torch.Tensor Preprocess input. Intended to be overwritten by subclass, if preprocessing is necessary. :param x [torch.Tensor]: Input tensor to preprocess. :returns: Processed tensor. .. py:method:: _alive(x) .. py:method:: _update(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. :returns: Residual cell update, BCWH. .. py:method:: _forward_step(x: torch.Tensor, step: int) .. py:method:: forward(x: torch.Tensor, steps: int = 1) -> ncalab.prediction.Prediction :param x [torch.Tensor]: Input image, padded along the channel dimension, BCWH. :param steps [int]: Time steps in forward pass. :returns [Prediction]: Prediction object. .. py:method:: _post_forward_step(x: torch.Tensor) -> torch.Tensor .. py:method:: 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. :param image [torch.Tensor]: Input image, BCWH. :param label [torch.Tensor]: Ground truth, BCWH. :returns: Dictionary of identifiers mapped to computed losses. .. py:method:: 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 .. py:method:: predict(image: torch.Tensor, steps: Optional[int | Tuple[int, int]] = None) -> ncalab.prediction.Prediction Make an NCA prediction, performing multiple forward passes to yield a final result. :param image: Input image, BCWH. :type image: torch.Tensor :param steps: Time steps :type steps: Optional[int] :returns: Prediction object. :rtype: Prediction .. py:method:: record(image: torch.Tensor, steps: Optional[int | Tuple[int, int]] = None) -> List[ncalab.prediction.Prediction] Record predictions for all time steps and return the resulting sequence of predictions. :param image: Input image, BCWH. :type image: torch.Tensor :returns: List of Prediction objects. :rtype: List[Prediction] .. py:method:: validate(dataloader: torch.utils.data.DataLoader, steps: Optional[int] = 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. :param dataloader [torch.utils.data.DataLoader]: Dataloader for validation images :param steps [int]: Inference steps :returns [Tuple[float, List[Prediction]]]: Validation metric, predicted image BCWH .. py:method:: _to_dict() -> Dict[str, Any] .. py:method:: to_dict() -> Dict[str, Any] .. py:method:: from_dict(d: Dict[str, Any]) :classmethod: .. py:method:: num_trainable_parameters() -> int Returns the number of trainable model parameters. :return: Number of trainable parameters. :rtype: int .. py:method:: save(path: str | os.PathLike) .. py:method:: load(model: AbstractNCAModel, path: str | os.PathLike) -> AbstractNCAModel :staticmethod: .. py:method:: post_prediction(prediction: ncalab.prediction.Prediction) -> ncalab.prediction.Prediction .. py:class:: AbstractNCAHead Bases: :py:obj:`torch.nn.Module`, :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: optimizer :value: None .. py:method:: forward(x: torch.Tensor) -> torch.Tensor :param x: Input tensor :type x: torch.Tensor :returns: NotImplemented, subclasses are required to implement this method. .. py:method:: freeze(freeze_last: bool = True) -> None :abstractmethod: Freeze head weights. :param freeze_last: Whether to freeze the last layer (if applicable), defaults to True :type freeze_last: bool, optional :returns: NotImplemented, subclasses are required to implement this method. .. py:class:: BasicNCAPerception(nca: ncalab.models.basicNCA.AbstractNCAModel) .. py:attribute:: nca .. py:method:: _define_filters() Define list of perception filters, based on parameters passed in constructor. :param num_learned_filters [int]: Number of learned filters in perception filter bank. .. py:method:: perceive(x: torch.Tensor, step: int) -> torch.Tensor .. py:method:: freeze() .. py:class:: AbstractNCARule(device: torch.device, input_size: int, hidden_size: int, output_size: int) Bases: :py:obj:`torch.nn.Module`, :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: device .. py:attribute:: input_size .. py:attribute:: hidden_size .. py:attribute:: output_size .. py:method:: freeze(freeze_last: bool = False) -> None :abstractmethod: .. py:class:: MLPNCARule(device: torch.device, input_size: int, hidden_size: int, output_size: int, nonlinearity: type[torch.nn.Module] = nn.ReLU) Bases: :py:obj:`ncalab.models.basicNCA.abstractNCArule.AbstractNCARule` NCA rule module based on a two-layer Multi-Layer-Perceptron (MLP). :param nn: _description_ :type nn: _type_ :param device: Compute device :type device: torch.device :param input_size: Input neurons :type input_size: int :param hidden_size: Hidden neurons :type hidden_size: int :param output_size: Output neurons :type output_size: int :param nonlinearity: Activation function, defaults to nn.ReLU :type nonlinearity: type[nn.Module], optional .. py:attribute:: nonlinearity .. py:method:: _build_network() .. py:method:: _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. .. py:method:: forward(x: torch.Tensor) -> torch.Tensor :param x: BCWH perception vector :type x: torch.Tensor :return: BCWH residual update :rtype: torch.Tensor .. py:method:: freeze(freeze_last: bool = False) Freeze the first layer of the NCA rule network and, optionally, the final layer. :param freeze_last: _description_, defaults to False :type freeze_last: bool, optional .. py:class:: CascadeNCA(wrapped: ncalab.models.basicNCA.AbstractNCAModel, scales: List[int], steps: List[int], single_model: bool = True) Bases: :py:obj:`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. :param wrapped: Backbone model based on AbstractNCAModel. :type wrapped: ncalab.AbstractNCAModel :param scales: List of scales to operate at, e.g. [4, 2, 1]. :type scales: List[int] :param steps: List of number of NCA inference time steps. :type steps: List[int] :param single_model: Only train a single instance of the NCA model :type single_model: bool .. py:attribute:: 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. :param image [torch.Tensor]: Input image, BCWH. :param label [torch.Tensor]: Ground truth, BCWH. :returns: Dictionary of identifiers mapped to computed losses. .. py:attribute:: finetune Prepare model for fine tuning by freezing everything except the final layer, and setting to "train" mode. :param: freeze_head .. py:attribute:: prepare_input Preprocess input. Intended to be overwritten by subclass, if preprocessing is necessary. :param x [torch.Tensor]: Input tensor to preprocess. :returns: Processed tensor. .. py:attribute:: head .. py:attribute:: metrics .. py:attribute:: wrapped .. py:attribute:: scales .. py:attribute:: steps .. py:attribute:: single_model :value: True .. py:attribute:: models :type: List[ncalab.models.basicNCA.AbstractNCAModel] .. py:method:: forward(x: torch.Tensor, *args, **kwargs) -> ncalab.prediction.Prediction :param x: Input image tensor, BCWH. :type x: torch.Tensor :param steps: Unused, as steps are defined in constructor. :type steps: torch.Tensor :return: Prediction object :rtype: Prediction .. py:method:: record(image: torch.Tensor, steps: Optional[int | Tuple[int, int]] = 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. :param image: Input image, BCWH. :type image: torch.Tensor :returns: List of Prediction objects. :rtype: List[Prediction] .. py:class:: 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: Optional[List[str]] = None, avg_pool_size: int = 8, lambda_hidden: float = 0, **kwargs) Bases: :py:obj:`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. :param device: Pytorch device descriptor. :param num_image_channels: _description_ :param num_hidden_channels: _description_ :param num_classes: _description_ :param fire_rate: Fire rate for stochastic weight update. Defaults to 0.8. :param hidden_size: Number of neurons in hidden layer. Defaults to 128. :param use_alive_mask: Whether to use alive masking (channel 3) during training. Defaults to False. :param pixel_wise_loss: Whether a prediction per pixel is desired, like in self-classifying MNIST. Defaults to False. :param num_learned_filters: Number of learned filters. If zero, use two sobel filters instead. Defaults to 2. :param filter_padding: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular". :param pad_noise: Whether to pad input image tensor with noise in hidden / output channels .. py:attribute:: _num_classes .. py:attribute:: pixel_wise_loss :value: False .. py:attribute:: use_classifier :value: True .. py:attribute:: avg_pool_size :value: 8 .. py:attribute:: lambda_hidden :value: 0 .. py:attribute:: metrics .. py:attribute:: focal_loss .. py:property:: num_classes :type: int .. py:method:: classify(image: torch.Tensor, steps: int = 100, reduce: bool = False) -> torch.Tensor Predict classification for an input image. :param image: Input image. :param steps: Inference steps. Defaults to 100. :param reduce: Return a single softmax probability. Defaults to False. :returns: Single class index or vector of logits. .. py:method:: loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) -> Dict[str, torch.Tensor] Return the classification loss. :param pred: Prediction. :param label: Ground truth. :returns: Dictionary of identifiers mapped to computed losses. .. py:method:: post_prediction(prediction: ncalab.prediction.Prediction) -> ncalab.prediction.Prediction .. py:class:: ClassificationNCAHead(num_hidden_channels: int, num_classes: int, device: torch.device, avg_pool_size: int, hidden_size: int = 32) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initialize internal Module state, shared by both nn.Module and ScriptModule. .. py:attribute:: num_hidden_channels .. py:attribute:: num_classes .. py:attribute:: device .. py:attribute:: avg_pool_size .. py:attribute:: classifier .. py:method:: forward(x) :param x: Input tensor :type x: torch.Tensor :returns: NotImplemented, subclasses are required to implement this method. .. py:method:: freeze(freeze_last: bool = False) Freeze head weights. :param freeze_last: Whether to freeze the last layer (if applicable), defaults to True :type freeze_last: bool, optional :returns: NotImplemented, subclasses are required to implement this method. .. py:class:: 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: :py:obj:`ncalab.models.basicNCA.AbstractNCAModel` NCA model for monocular depth estimation. :param device: Pytorch device descriptor. :param num_image_channels: Number of channels reserved for input image. :param num_hidden_channels: Number of hidden channels (communication channels). :param num_output_channels: Number of output channels. :param validation_metric: :param fire_rate: Fire rate for stochastic weight update. Defaults to 0.5. :param hidden_size: Number of neurons in hidden layer. Defaults to 128. :param use_alive_mask: Whether to use alive masking (channel 3) during training. Defaults to False. :param immutable_image_channels: If image channels should be fixed during inference, which is the case for most segmentation or classification problems. Defaults to True. :param num_learned_filters: Number of learned filters. If zero, use two sobel filters instead. Defaults to 2. :param filter_padding: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular". :param use_laplace: Whether to use Laplace filter (only if num_learned_filters == 0) :param kernel_size: Filter kernel size (only for learned filters) :param pad_noise: Whether to pad input image tensor with noise in hidden / output channels :param use_temporal_encoding: :param rule_type: :param rule_args: :param training_timesteps: :param inference_timesteps: .. py:attribute:: vignette :value: None .. py:method:: loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) -> Dict[str, torch.Tensor] :param image: Input image, BCWH. :param label: Ground truth. :returns: Dictionary of identifiers mapped to computed losses. .. py:class:: 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: :py:obj:`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. :param device [torch.device]: Pytorch device descriptor. :param num_image_channels [int]: Number of channels reserved for input image. Defaults to 4. :param num_hidden_channels [int]: Number of hidden channels (communication channels). Defaults to 16. :param fire_rate [float]: Stochastic weight update. Defaults to 0.5. :param hidden_size [int]: Default number of nodes in hidden layer. Defaults to 128. :param use_alive_mask [bool]: Whether to use alive masking. Defaults to False. .. py:attribute:: lambda_hidden :value: 0.0 .. py:method:: loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) -> Dict[str, torch.Tensor] Implements a simple MSE loss between target and prediction. :param pred: Prediction :param label: Target :returns [Tensor]: MSE Loss .. py:method:: make_seed(width: int, height: int) -> torch.Tensor .. py:method:: grow(seed: torch.Tensor, steps: int = 100) -> List[numpy.ndarray] Run the growth process and return the resulting output sequence. :param seed [torch.Tensor]: Seed image, can be generated through make_seed. :param steps [int]: Number of inference steps. Defaults to 100. :returns [List[np.ndarray]]: Sequence of output images. .. py:class:: 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: :py:obj:`ncalab.models.basicNCA.AbstractNCAModel` Model used for image segmentation. Uses Dice score as the default validation metric. Currently, only binary segmentation masks are supported. :param device [torch.device]: Compute device. :param num_image_channels [int]: Number of image channels. Defaults to 3. :param num_hidden_channels [int]: Number of hidden channels. Defaults to 16. :param num_classes [int]: Number of classes. Defaults to 1. :param fire_rate [float]: NCA fire rate. Defaults to 0.8. :param hidden_size [int]: Number of neurons in hidden layer. Defaults to 128. :param learned_filters [int]: Number of learned filters. If 0, use sobel. Defaults to 2. :param pad_noise [bool]: Whether to pad input images with noise. Defaults to True. :param filter_padding [str]: Padding type to use. Might affect reliance on spatial cues. Defaults to "circular". .. py:attribute:: num_classes :value: 1 .. py:attribute:: metrics .. py:attribute:: lambda_hidden :value: 0.001 .. py:attribute:: bce_loss .. py:attribute:: dice_loss .. py:method:: loss(pred: ncalab.prediction.Prediction, label: torch.Tensor) -> Dict[str, torch.Tensor] Compute Dice loss. :param pred: Prediction. :param label: Ground truth. :returns: Dictionary of identifiers mapped to computed losses. .. py:method:: _post_forward_step(x: torch.Tensor) -> torch.Tensor .. py:method:: post_prediction(prediction: ncalab.prediction.Prediction) -> ncalab.prediction.Prediction