ncalab.models.basicNCA.mlpNCArule

Classes

MLPNCARule

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

Module Contents

class ncalab.models.basicNCA.mlpNCArule.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