nn

The nn module, a key part of the minima’s toolkit, is a pretty cool and all-encompassing software.

source

Parameter

 Parameter (array, device:Optional[minima.autograd.Device]=None,
            dtype=None, requires_grad=True, **kwargs)

A kind of Tensor that is to be considered a module parameter.

Parameters are Tensor subclasses, that have a very special property when used with Module s - when they’re assigned as Module attributes they are automatically added to the list of its parameters, and will appear in Module.parameters() iterator. Another difference is that parameters can’t be volatile and that they require gradient by default.

source

Module

 Module ()

Base class for all neural network modules in Minima.

Your models should also subclass this class. Subclasses should define a forward method.

Attributes: - training (bool): Module is initialized in training mode by default. Use eval() to switch it to evaluation mode.

Methods: - parameters(): Returns a list of all Parameter instances in the module. - _children(): Returns a list of all child Module instances. - eval(): Switches the module and all its children to evaluation mode. - train(): Switches the module and all its children back to training mode. - __call__(): The call method, which simply calls the forward method, must be defined by all subclasses.

source

Sequential

 Sequential (*modules)

A sequential container in Minima.

Modules will be added to it in the order they are passed in the constructor. A Sequential module contains a sequence of child modules stored in the order they were added. Each module is applied in order to the input to produce the output.

The Sequential class makes it easy to build networks where the output of one layer is the input to the next.

Attributes: - modules (tuple of Module): The sequence of child modules to apply.

Methods: - forward(x: Tensor) -> Tensor: Passes the input through all the child modules in sequential order.

source

Linear

 Linear (in_features, out_features, bias=True, device=None,
         dtype='float32')

A class representing a fully connected (linear) layer in a neural network. This class inherits from the Module class.

Attributes: in_features (int): The number of input features. out_features (int): The number of output features. device (str): The device to store the Parameters on (defaults to None, which means CPU). dtype (str): The data type of the Parameters (defaults to ‘float32’). weight (Parameter): The weight parameters of the layer. bias (Parameter): The bias parameters of the layer, or None if bias=False.

Methods: forward(X: Tensor) -> Tensor: Compute the forward pass of the layer.

Type Default Details
in_features The number of input features.
out_features The number of output features.
bias bool True Whether or not to include a bias term. Default is True.
device NoneType None The device to store the Parameters on. Default is None, which means CPU.
dtype str float32 The data type of the Parameters. Default is ‘float32’. 
from minima.optim import *
class MyModule(Module):
    def __init__(self):
        super().__init__()
        self.layer1 = mi.nn.Linear(10, 20)
        self.layer2 = mi.nn.Linear(20, 10)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x

my_module = MyModule()
my_module
MyModule()
my_module.parameters()
[]
tt = mi.nn.Linear(in_features=10, out_features=5, bias=True)
isinstance(tt, mi.nn.Module)
True

source

Flatten

 Flatten ()

A Flatten module in Minima.

This module flattens an input tensor into a 2D matrix, typically for transitioning from convolutional layers to linear layers within a neural network model.

Methods: - forward(X: Tensor) -> Tensor: Flattens the input tensor.

source

ReLU

 ReLU ()

Base class for all neural network modules in Minima.

Your models should also subclass this class. Subclasses should define a forward method.

Attributes: - training (bool): Module is initialized in training mode by default. Use eval() to switch it to evaluation mode.

Methods: - parameters(): Returns a list of all Parameter instances in the module. - _children(): Returns a list of all child Module instances. - eval(): Switches the module and all its children to evaluation mode. - train(): Switches the module and all its children back to training mode. - __call__(): The call method, which simply calls the forward method, must be defined by all subclasses.

source

Sigmoid

 Sigmoid ()

Base class for all neural network modules in Minima.

Your models should also subclass this class. Subclasses should define a forward method.

Attributes: - training (bool): Module is initialized in training mode by default. Use eval() to switch it to evaluation mode.

Methods: - parameters(): Returns a list of all Parameter instances in the module. - _children(): Returns a list of all child Module instances. - eval(): Switches the module and all its children to evaluation mode. - train(): Switches the module and all its children back to training mode. - __call__(): The call method, which simply calls the forward method, must be defined by all subclasses.

The implementation you’ve shared is a numerically stable version of the Cross Entropy Loss formula, which is generally defined for a single sample as:

\[ H(p, q) = - \sum_{i} p_i \log(q_i) \]

where: - \(p\) is the true distribution (in classification, typically a one-hot encoded vector), - \(q\) is the predicted distribution (output of the softmax function on the logits from the model).

However, directly implementing this formula can lead to numerical instability because of the log operation. The given implementation overcomes this by using the Log-Sum-Exp trick to prevent underflow or overflow. \[ CE = -\sum_{i=1}^{C} y_i \log(\hat{y_i}) \]

where \(y\) is the ground truth label and \(\hat{y}\) is the predicted probability, and \(C\) is the number of classes. In the case of one-hot encoding, only the term corresponding to the true class contributes to the sum. So, we can simplify it to:

\[ CE = -\log(\hat{y_c}) \]

where \(c\) is the correct class.

The predicted probabilities \(\hat{y}\) are obtained by applying the softmax function to the logits \(z\):

\[ \hat{y_i} = \frac{e^{z_i}}{\sum_{j=1}^{C} e^{z_j}} \]

Substituting \(\hat{y_c}\) in the Cross-Entropy Loss, we have:

\[ CE = -\log\left(\frac{e^{z_c}}{\sum_{j=1}^{C} e^{z_j}}\right) \]

Applying the logarithm property \(\log(a/b) = \log(a) - \log(b)\), we get:

\[ CE = -z_c + \log\left(\sum_{j=1}^{C} e^{z_j}\right) \]

First, log_sum_exp_logits = operators.logsumexp(logits, axes=(1, )).sum() computes the term \(\log\left(\sum_{j=1}^{C} e^{z_j}\right)\). The function logsumexp computes the logarithm of the sum of exponentials in a numerically stable way, and then these values are summed over all samples.

Second, true_class_logits_sum = (logits * init.one_hot(logits.shape[1], y)).sum() computes the \(-z_c\) term for each sample. The function init.one_hot(logits.shape[1], y) creates a one-hot encoding of the true labels, and this is then multiplied with the logits to pick out the logits for the correct classes. These values are then summed over all samples.

Finally, (log_sum_exp_logits - true_class_logits_sum) / logits.shape[0] computes the average loss per sample.

logits = mi.Tensor([[ 0.6734,  0.2576],
        [ 0.4689,  0.4607],
        [-2.2457, -0.3727],
        [ 4.4164, -1.2760],
        [ 0.9233,  0.5347],
        [ 1.0698,  1.6187]])
targ = mi.Tensor([0,1,0,1,1,0])
log_sum_exp_logits = operators.logsumexp(logits, axes=(1, ))
log_sum_exp_logits
minima.Tensor(
[ 1.180104  1.157956 -0.229759  4.419766  1.440906  2.074595])
log_sum_exp_logits = log_sum_exp_logits.sum()
one_hot_y = init.one_hot(logits.shape[1], targ)
one_hot_y
minima.Tensor(
[[1. 0.]
 [0. 1.]
 [1. 0.]
 [0. 1.]
 [0. 1.]
 [1. 0.]])
true_class_logits = (logits * one_hot_y)
true_class_logits
minima.Tensor(
[[ 0.6734  0.    ]
 [ 0.      0.4607]
 [-2.2457 -0.    ]
 [ 0.     -1.276 ]
 [ 0.      0.5347]
 [ 1.0698  0.    ]])
true_class_logits_sum = true_class_logits.sum()
true_class_logits_sum
minima.Tensor(
-0.7830999999999999)
loss = (log_sum_exp_logits - true_class_logits_sum) / targ.shape[0]
loss
minima.Tensor(
1.8044446680186768)
loss = torch.nn.CrossEntropyLoss()(torch.tensor(logits.numpy()), torch.tensor(targ.numpy()))
loss
tensor(1.8044, dtype=torch.float64)

source

CrossEntropyLoss

 CrossEntropyLoss ()

Cross-entropy loss module in Minima.

This module computes the Cross Entropy Loss between the input logits and the target classes. It’s useful in classification tasks where the model outputs probabilities for each class.

Methods: - forward(input: Tensor, target: Tensor) -> Tensor: Calculates the cross-entropy loss between the input (logits) and the target (class indices).

Example:

model = Sequential(
    Linear(10, 20),
    ReLU(),
    Linear(20, 10),
)
loss_fn = CrossEntropyLoss()
output = model(input_tensor)  # compute model output
loss = loss_fn(output, target_tensor)  # compute loss

source

Softmax

 Softmax ()

Cross-entropy loss module in Minima.

This module computes the Cross Entropy Loss between the input logits and the target classes. It’s useful in classification tasks where the model outputs probabilities for each class.

Methods: - forward(input: Tensor, target: Tensor) -> Tensor: Calculates the cross-entropy loss between the input (logits) and the target (class indices).

Example:

model = Sequential(
    Linear(10, 20),
    ReLU(),
    Linear(20, 10),
)
loss_fn = CrossEntropyLoss()
output = model(input_tensor)  # compute model output
loss = loss_fn(output, target_tensor)  # compute loss
sm = Softmax()
t = mi.Tensor([[ 1.25721,   1.066163],
               [ 1.123695, -0.410479],
               [ 0.047494, -1.200239],
               [ 0.440039,  0.884834],])
sm(t)
minima.Tensor(
[[0.547617 0.452383]
 [0.822616 0.177384]
 [0.776907 0.223093]
 [0.390599 0.609401]])

Layer normalization

Layer normalization \(\text{LN}\) normalizes the input \(X\) as follows:

When input \(X \in \mathbb{R}^{B \times C}\) is a batch of embeddings, where \(B\) is the batch size and \(C\) is the number of features. \(\gamma \in \mathbb{R}^{C}\) and \(\beta \in \mathbb{R}^{C}\). \[\text{LN}(X) = \gamma \frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}} + \beta\]

When input \(X \in \mathbb{R}^{L \times B \times C}\) is a batch of a sequence of embeddings, where \(B\) is the batch size, \(C\) is the number of channels, \(L\) is the length of the sequence. \(\gamma \in \mathbb{R}^{C}\) and \(\beta \in \mathbb{R}^{C}\). \[\text{LN}(X) = \gamma \frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}} + \beta\]

When input \(X \in \mathbb{R}^{B \times C \times H \times W}\) is a batch of image representations, where \(B\) is the batch size, \(C\) is the number of channels, \(H\) is the height and \(W\) is the width. This is not a widely used scenario. \(\gamma \in \mathbb{R}^{C \times H \times W}\) and \(\beta \in \mathbb{R}^{C \times H \times W}\). \[\text{LN}(X) = \gamma \frac{X - \underset{C, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{C, H, W}{Var}[X] + \epsilon}} + \beta\]

X = mi.Tensor(init.rand(5, 10))
X
minima.Tensor(
[[0.424498 0.53646  0.553287 0.97566  0.097845 0.951467 0.582025 0.259501 0.248237 0.651587]
 [0.762057 0.323642 0.67996  0.59344  0.998091 0.910109 0.327038 0.583429 0.118361 0.912831]
 [0.503303 0.953529 0.739612 0.116599 0.830379 0.703346 0.34279  0.722535 0.687944 0.927093]
 [0.874576 0.247239 0.280458 0.342372 0.324733 0.963528 0.838563 0.806436 0.37879  0.214825]
 [0.659279 0.326977 0.308499 0.127876 0.815616 0.123615 0.956371 0.24132  0.681585 0.669349]])
bs, fs = X.shape
bs, fs
(5, 10)
mean = X.sum(axes=(1,)) / fs
mean
minima.Tensor(
[0.528057 0.620896 0.652713 0.527152 0.491049])
mean = mean.reshape((bs, 1))
mean
minima.Tensor(
[[0.528057]
 [0.620896]
 [0.652713]
 [0.527152]
 [0.491049]])
mean = mean.broadcast_to(X.shape)
mean
minima.Tensor(
[[0.528057 0.528057 0.528057 0.528057 0.528057 0.528057 0.528057 0.528057 0.528057 0.528057]
 [0.620896 0.620896 0.620896 0.620896 0.620896 0.620896 0.620896 0.620896 0.620896 0.620896]
 [0.652713 0.652713 0.652713 0.652713 0.652713 0.652713 0.652713 0.652713 0.652713 0.652713]
 [0.527152 0.527152 0.527152 0.527152 0.527152 0.527152 0.527152 0.527152 0.527152 0.527152]
 [0.491049 0.491049 0.491049 0.491049 0.491049 0.491049 0.491049 0.491049 0.491049 0.491049]])
x_centred = X - mean
x_centred
minima.Tensor(
[[-0.103559  0.008403  0.02523   0.447604 -0.430211  0.42341   0.053968 -0.268556 -0.27982   0.12353 ]
 [ 0.141161 -0.297254  0.059064 -0.027456  0.377195  0.289213 -0.293858 -0.037467 -0.502535  0.291935]
 [-0.14941   0.300816  0.086899 -0.536114  0.177666  0.050634 -0.309923  0.069822  0.035231  0.27438 ]
 [ 0.347424 -0.279913 -0.246694 -0.18478  -0.202419  0.436376  0.311411  0.279284 -0.148362 -0.312327]
 [ 0.16823  -0.164072 -0.18255  -0.363173  0.324568 -0.367433  0.465322 -0.249728  0.190536  0.1783  ]])

source

LayerNorm1d

 LayerNorm1d (dim:int, eps=1e-05, device=None, dtype='float32')

1D Layer normalization module in Minima.

Applies layer normalization over a 1D input. The mean and standard deviation are computed over the last dimension.

Attributes: - dim (int): The dimension of the input feature space. - eps (float): A small constant for numerical stability. - weight (Parameter): The learnable weights of the module of size ‘dim’, initialized with ones. - bias (Parameter): The learnable bias of the module of size ‘dim’, initialized with zeros.

Methods: - forward(x: Tensor) -> Tensor: Applies layer normalization to the input tensor.

Type Default Details
dim int The dimension of the input feature space.
eps float 1e-05 A small constant for numerical stability. Default is 1e-5.
device NoneType None The desired device of returned tensor. If None, uses the current device for the default tensor type. Default is None.
dtype str float32 The desired data type of returned tensor. If None, uses the default data type. Default is “float32”.

Batch Norm

This is an implementation of Batch Normalization from paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.

Internal Covariate Shift

The paper defines Internal Covariate Shift as the change in the distribution of network activations due to the change in network parameters during training. For example, let’s say there are two layers \(l_1\) and \(l_2\). During the beginning of the training \(l_1\) outputs (inputs to \(l_2\)) could be in distribution \(\mathcal{N}(0.5, 1)\). Then, after some training steps, it could move to \(\mathcal{N}(0.6, 1.5)\). This is internal covariate shift.

Internal covariate shift will adversely affect training speed because the later layers (\(l_2\) in the above example) have to adapt to this shifted distribution.

By stabilizing the distribution, batch normalization minimizes the internal covariate shift.

Normalization

It is known that whitening improves training speed and convergence. Whitening is linearly transforming inputs to have zero mean, unit variance, and be uncorrelated.

Normalizing outside gradient computation doesn’t work

Normalizing outside the gradient computation using pre-computed (detached) means and variances doesn’t work. For instance. (ignoring variance), let \[\hat{x} = x - \mathbb{E}[x]\] where \(x = u + b\) and \(b\) is a trained bias and \(\mathbb{E}[x]\) is an outside gradient computation (pre-computed constant).

Note that \(\hat{x}\) has no effect on \(b\). Therefore, \(b\) will increase or decrease based \(\frac{\partial{\mathcal{L}}}{\partial x}\), and keep on growing indefinitely in each training update. The paper notes that similar explosions happen with variances.

Batch Normalization

Whitening is computationally expensive because you need to de-correlate and the gradients must flow through the full whitening calculation.

The paper introduces a simplified version which they call Batch Normalization. First simplification is that it normalizes each feature independently to have zero mean and unit variance: \[\hat{x}^{(k)} = \frac{x^{(k)} - \mathbb{E}[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}\] where \(x = (x^{(1)} ... x^{(d)})\) is the \(d\)-dimensional input.

The second simplification is to use estimates of mean \(\mathbb{E}[x^{(k)}]\) and variance \(Var[x^{(k)}]\) from the mini-batch for normalization; instead of calculating the mean and variance across the whole dataset.

Normalizing each feature to zero mean and unit variance could affect what the layer can represent. As an example paper illustrates that, if the inputs to a sigmoid are normalized most of it will be within \([-1, 1]\) range where the sigmoid is linear. To overcome this each feature is scaled and shifted by two trained parameters \(\gamma^{(k)}\) and \(\beta^{(k)}\). \[y^{(k)} =\gamma^{(k)} \hat{x}^{(k)} + \beta^{(k)}\] where \(y^{(k)}\) is the output of the batch normalization layer.

Note that when applying batch normalization after a linear transform like \(Wu + b\) the bias parameter \(b\) gets cancelled due to normalization. So you can and should omit bias parameter in linear transforms right before the batch normalization.

Batch normalization also makes the back propagation invariant to the scale of the weights and empirically it improves generalization, so it has regularization effects too.

Inference

We need to know \(\mathbb{E}[x^{(k)}]\) and \(Var[x^{(k)}]\) in order to perform the normalization. So during inference, you either need to go through the whole (or part of) dataset and find the mean and variance, or you can use an estimate calculated during training. The usual practice is to calculate an exponential moving average of mean and variance during the training phase and use that for inference.

Batch normalization layer \(\text{BN}\) normalizes the input \(X\) as follows:

When input \(X \in \mathbb{R}^{B \times C \times H \times W}\) is a batch of image representations, where \(B\) is the batch size, \(C\) is the number of channels, \(H\) is the height and \(W\) is the width. \(\gamma \in \mathbb{R}^{C}\) and \(\beta \in \mathbb{R}^{C}\). \[\text{BN}(X) = \gamma \frac{X - \underset{B, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{B, H, W}{Var}[X] + \epsilon}} + \beta\]

When input \(X \in \mathbb{R}^{B \times C}\) is a batch of embeddings, where \(B\) is the batch size and \(C\) is the number of features. \(\gamma \in \mathbb{R}^{C}\) and \(\beta \in \mathbb{R}^{C}\). \[\text{BN}(X) = \gamma \frac{X - \underset{B}{\mathbb{E}}[X]}{\sqrt{\underset{B}{Var}[X] + \epsilon}} + \beta\]

When input \(X \in \mathbb{R}^{B \times C \times L}\) is a batch of a sequence embeddings, where \(B\) is the batch size, \(C\) is the number of features, and \(L\) is the length of the sequence. \(\gamma \in \mathbb{R}^{C}\) and \(\beta \in \mathbb{R}^{C}\). \[\text{BN}(X) = \gamma \frac{X - \underset{B, L}{\mathbb{E}}[X]}{\sqrt{\underset{B, L}{Var}[X] + \epsilon}} + \beta\]

source

BatchNorm1d

 BatchNorm1d (dim:int, eps=1e-05, momentum=0.1, device=None,
              dtype='float32')

1D Batch normalization module in Minima.

This module applies Batch Normalization over a 1D input as described in the paper “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift” by Ioffe and Szegedy.

Attributes: - dim (int): The dimension of the input feature space. - eps (float): A small constant added to the denominator for numerical stability. - momentum (float): The value used for the running_mean and running_var computation. - weight (Parameter): The learnable scale factor of the module of size ‘dim’, initialized with ones. - bias (Parameter): The learnable offset of the module of size ‘dim’, initialized with zeros. - running_mean (Tensor): The running mean. Represents the mean of the features over batches. Initialized with zeros. - running_std (Tensor): The running standard deviation. Represents the standard deviation of the features over batches. Initialized with ones.

Methods: - update_stats(x: Tensor) -> Tuple[Tensor, Tensor]: Calculates the mean and standard deviation of the input tensor. - forward(x: Tensor) -> Tensor: Applies batch normalization to the input tensor.

Example:

batch_norm = BatchNorm1d(dim=512)
output = batch_norm(input_tensor)  # Apply batch normalization
/opt/hostedtoolcache/Python/3.9.17/x64/lib/python3.9/site-packages/fastcore/docscrape.py:225: UserWarning: potentially wrong underline length... 
Parameters: 
---------- in 
Dropout Layer for a Neural Network.
...
  else: warn(msg)

source

Dropout

 Dropout (p=0.5)

Dropout Layer for a Neural Network.

This class represents a dropout layer in a neural network, which is a simple and effective regularization technique. During training, it randomly zeroes out some of the elements of the input tensor with probability p using samples from a Bernoulli distribution.

Parameters:

p: float, optional, default = 0.5 Probability of an element to be zeroed. Default: 0.5.

/opt/hostedtoolcache/Python/3.9.17/x64/lib/python3.9/site-packages/fastcore/docscrape.py:225: UserWarning: potentially wrong underline length... 
Parameters: 
---------- in 
Residual Layer for a Neural Network.
...
  else: warn(msg)

source

Residual

 Residual (fn:__main__.Module)

Residual Layer for a Neural Network.

This class represents a residual layer in a neural network, which is a technique that helps to overcome the problem of vanishing and exploding gradients in deep neural networks. It achieves this by allowing gradients to pass through layers directly (via an identity shortcut connection) without any modification.

Parameters:

fn: Module The function to be applied to the input tensor.

source

Identity

 Identity ()

Base class for all neural network modules in Minima.

Your models should also subclass this class. Subclasses should define a forward method.

Attributes: - training (bool): Module is initialized in training mode by default. Use eval() to switch it to evaluation mode.

Methods: - parameters(): Returns a list of all Parameter instances in the module. - _children(): Returns a list of all child Module instances. - eval(): Switches the module and all its children to evaluation mode. - train(): Switches the module and all its children back to training mode. - __call__(): The call method, which simply calls the forward method, must be defined by all subclasses.

Export

import nbdev; nbdev.nbdev_export()