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