ncalab.models.basicNCA ====================== .. py:module:: ncalab.models.basicNCA Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/ncalab/models/basicNCA/abstractNCA/index /autoapi/ncalab/models/basicNCA/abstractNCAhead/index /autoapi/ncalab/models/basicNCA/abstractNCArule/index /autoapi/ncalab/models/basicNCA/basicNCAperception/index /autoapi/ncalab/models/basicNCA/mlpNCArule/index Classes ------- .. autoapisummary:: ncalab.models.basicNCA.AbstractNCAModel ncalab.models.basicNCA.AbstractNCAHead ncalab.models.basicNCA.BasicNCAPerception ncalab.models.basicNCA.AbstractNCARule ncalab.models.basicNCA.MLPNCARule 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