yamle.models.operations module

yamle.models.operations.output_activation(x, task, dim=None)

This function applies the output activation.

Parameters:

• x (torch.Tensor) – The input tensor.

• task (str) – If the task is ‘classification’, the output is ‘softmax’. If the task is ‘regression’, the output is ‘exp()’ for the variance if the output is of shape (batch_size, 2).

• dim (Optional[int]) – The dimension to apply the activation on. Defaults to 1.

Return type:

Tensor

class yamle.models.operations.OutputActivation(task, dim=None)

Bases: Module

This class is used to apply the output activation.

If the task is classification, the output is softmax. If the task is regression, the output is exp() for the variance if the output is of shape (batch_size, 2).

Parameters:

• task (str) – The task to perform.

• dim (Optional[int]) – The dimension to apply the activation on. Defaults to 1.

forward(x)

This function applies the output activation.

Return type:

Tensor

extra_repr()

This function returns the extra representation of the output activation.

Return type:

str

enable()

Enable the activation.

Return type:

None

disable()

Disable the activation.

Return type:

None

training: bool

class yamle.models.operations.ReshapeOutput(num_members)

Bases: Module

This class is used to reshape the output of the model depending on the number of members.

It does so with respect to the second dimension that is created by the num_members from the third dimension.

forward(x)

This function reshapes the output of the model.

Return type:

Tensor

extra_repr()

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

training: bool

class yamle.models.operations.ReshapeInput(*args, **kwargs)

Bases: Module

This class is used to reshape the input of the model depending on the number of members.

It folds the num_members dimension into the second dimension.

forward(x)

This function reshapes the input of the model.

Return type:

Tensor

training: bool

class yamle.models.operations.Unsqueeze(shape_length)

Bases: Module

This class is used to unsqueeze a tensor to a given shape length.

forward(x)

This function is used to unsqueeze the tensor.

Return type:

Tensor

training: bool

class yamle.models.operations.Add(inplace=False)

Bases: Module

A simple class implementing residual addition.

The forward function is simply the addition of the two inputs.

Parameters:

inplace (bool) – If True, the addition is done in-place.

forward(x, y)

The forward function of the residual addition.

Return type:

Tensor

training: bool

class yamle.models.operations.Multiply(inplace=False)

Bases: Module

A simple class implementing residual multiplication.

The forward function is simply the multiplication of the two inputs.

Parameters:

inplace (bool) – If True, the multiplication is done in-place.

forward(x, y)

The forward function of the residual multiplication.

Return type:

Tensor

training: bool

class yamle.models.operations.ResidualLayer(layer, inplace=False)

Bases: Module

This class implements a residual layer.

It consists of a layer followed by a residual addition with the input. The layer should be a nn.Module of nn.Sequential type.

Parameters:

• layer (nn.Module) – The layer to be used.

• inplace (bool) – If True, the addition is done in-place.

forward(x, identity=None)

The forward function of the residual layer.

By default, the identity is the input x. It can be changed by passing a tensor to the identity argument.

Return type:

Tensor

training: bool

class yamle.models.operations.ParallelModel(models, single_source=False, inputs_dim=1, outputs_dim=1, initialise_members_same=False)

Bases: Module

This class implements a parallel model.

It consists of a list of models that are applied in parallel to the input. The models should be a list of nn.Module of nn.Sequential type. The input is assumed to be of shape (batch_size, len(models), *) or (batch_size, *) if single_source is True.

Parameters:

• models (List[nn.Module]) – The models to be used.

• single_source (bool) – If True, the input is assumed to be of shape (batch_size, *). If False, the input is assumed to be of shape (batch_size, len(models), *). Defaults to False.

• inputs_dim (int) – The dimension to split the input on. Defaults to 1. Used only if single_source is False.

• outputs_dim (int) – The dimension to stack the outputs on. Defaults to 1.

• initialise_members_same (bool) – If True, the members of the models are initialised to the same values. Defaults to False.

forward(x)

The forward function of the parallel model.

Return type:

Tensor

initialise_members_same()

This is a helper function to initialise the members of the parallel model with the same weights.

The weights are copied from the first model in the list to all the other models.

Return type:

None

training: bool

class yamle.models.operations.Normalization(norm=None, dimension=1, norm_kwargs={})

Bases: Module

This class implements a normalization layer.

Parameters:

• norm (Optional[str]) – The type of normalization to use. Defaults to None. Choices are batch, layer and instance.

• dimension (int) – The dimension to normalise on. Defaults to 1.

• norm_kwargs (Dict[str, Any]) – The keyword arguments to be passed to the normalization layer. Defaults to {}.

forward(x)

The forward function of the normalization layer.

Return type:

Tensor

training: bool

class yamle.models.operations.Pooling(pooling=None, dimension=1, pool_kwargs={})

Bases: Module

This class implements a pooling layer.

Parameters:

• pooling (Optional[str]) – The type of pooling to use. Defaults to None. Choices are max, avg, adaptive_max and adaptive_avg.

• dimension (int) – The dimension to pool on. Defaults to 1.

• pool_kwargs (Dict[str, Any]) – The keyword arguments to be passed to the pooling layer. Defaults to {}.

forward(x)

The forward function of the pooling layer.

Return type:

Tensor

training: bool

class yamle.models.operations.Activation(activation=None, dimension=1)

Bases: Module

This class implements an activation function.

Parameters:

• activation (Optional[str]) – The type of activation to use. Defaults to None. Choices are relu, sigmoid, tanh, softmax and log_softmax.

• dimension (int) – The dimension to apply the activation on. Defaults to 1.

forward(x)

The forward function of the activation layer.

Return type:

Tensor

training: bool

class yamle.models.operations.MatrixMultiplication(*args, **kwargs)

Bases: Module

This class implements a matrix multiplication layer.

forward(x, y)

The forward function of the matrix multiplication layer.

Return type:

Tensor

training: bool

class yamle.models.operations.LinearNormActivation(in_features, out_features, bias=True, normalization=None, activation='relu')

Bases: Module

This class is used to create a linear layer followed by a normalization layer and an activation layer.

Parameters:

• in_features (int) – The number of input features.

• out_features (int) – The number of output features.

• bias (bool) – If True, the layer has a bias.

• normalization (Optional[str]) – The type of normalization to use. Defaults to None. Choices are batch, layer and instance.

• activation (str) – The type of activation to use. Defaults to relu. Choices are relu, sigmoid, tanh, softmax and log_softmax.

forward(x)

The forward function of the linear layer followed by normalization and ReLU.

Return type:

Tensor

replace_layers_for_quantization()

This method is used to replace the layers for quantization.

It merges the linear layer with the normalization layer and activation layer if possible. This is currently only possible if the normalization layer is batch normalization and the activation layer is ReLU.

Return type:

None

training: bool

class yamle.models.operations.Conv2dNormActivation(in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, bias=True, normalization=None, activation='relu')

Bases: Module

This class is used to create a convolutional layer followed by a normalization layer and an activation layer.

Parameters:

• in_channels (int) – The number of input channels.

• out_channels (int) – The number of output channels.

• kernel_size (int) – The size of the convolutional kernel. Default: 3.

• stride (int) – The stride of the convolution. Default: 1.

• padding (int) – The padding of the convolution. Default: 1.

• dilation (int) – The dilation of the convolution. Default: 1.

• groups (int) – The number of groups in the convolution. Default: 1.

• bias (bool) – If True, the convolution has a bias. Default: True.

• normalization (Optional[str]) – The type of normalization to use. Defaults to None. Choices are batch, layer and instance.

• activation (str) – The type of activation to use. Defaults to relu. Choices are relu, sigmoid, tanh, softmax and log_softmax.

forward(x)

The forward function of the convolutional norm ReLU layer.

Return type:

Tensor

replace_layers_for_quantization()

This method is used to replace the layers for quantization.

It merges the convolutional layer with the normalization layer and activation layer if possible. This is currently only possible if the normalization layer is batch normalization and the activation layer is ReLU.

Return type:

None

training: bool

class yamle.models.operations.DoubleConv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True, normalization='batch', activation='relu', residual=True)

Bases: Module

This class is used to create a double convolutional layer.

It is composed of two convolutional layers, followed by a normalization layer and ReLU. The first convolutional layer has output channels size out_channels. If residual is True, the output of the second convolutional layer is added to the input of the first convolutional layer.

Parameters:

• in_channels (int) – The number of input channels.

• out_channels (int) – The number of output channels.

• kernel_size (int) – The size of the convolutional kernel. Default: 3.

• stride (int) – The stride of the convolution. Default: 1.

• padding (int) – The padding of the convolution. Default: 1.

• bias (bool) – If True, the convolution has a bias. Default: True.

• normalization (Optional[str]) – The type of normalization to use. Defaults to batch. Choices are batch, layer and instance.

• activation (str) – The type of activation to use. Defaults to relu. Choices are relu, sigmoid, tanh, softmax and log_softmax.

• residual (bool) – If True, the output of the second convolutional layer is added to the input of the first convolutional layer. Default: True.

forward(x)

The forward function of the double convolutional layer.

Return type:

Tensor

training: bool

class yamle.models.operations.DepthwiseSeparableConv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True, normalization=None)

Bases: Module

This class is used to create a depthwise separable convolutional layer.

It consists of a depthwise convolutional layer followed by a pointwise convolutional layer.

Parameters:

• in_channels (int) – The number of input channels.

• out_channels (int) – The number of output channels.

• kernel_size (int) – The size of the convolutional kernel. Default: 3.

• stride (int) – The stride of the convolution. Default: 1.

• padding (int) – The padding of the convolution. Default: 1.

• bias (bool) – If True, the convolution has a bias. Default: True.

• normalization (Optional[str]) – The type of normalization to use. Defaults to None. Choices are batch, layer and instance.

reset_parameters()

This function is used to initialize the parameters of the layer.

Return type:

None

forward(x)

The forward function of the separable convolutional layer.

Return type:

Tensor

training: bool

class yamle.models.operations.CompletelySeparableConv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True, normalization=None)

Bases: Module

This class is used to create a completely separable convolutional layer.

It consists of two depthwise convolutional layers followed by a pointwise convolutional layer.

Parameters:

• in_channels (int) – The number of input channels.

• out_channels (int) – The number of output channels.

• kernel_size (int) – The size of the convolutional kernel. Default: 3.

• stride (int) – The stride of the convolution. Default: 1.

• padding (int) – The padding of the convolution. Default: 1.

• bias (bool) – If True, the convolution has a bias. Default: True.

• normalization (Optional[str]) – The type of normalization to use. Defaults to None. Choices are batch, layer and instance.

reset_parameters()

This function is used to initialize the parameters of the layer.

Return type:

None

forward(x)

The forward function of the separable convolutional layer.

Return type:

Tensor

training: bool

class yamle.models.operations.SqueezeAndExcitation(in_out_channels, reduction=16, activation='relu')

Bases: Module

This class is used to create a squeeze and excitation layer.

It is implemented according to the paper Squeeze-and-Excitation Networks.

Parameters:

• in_out_channels (int) – The number of input channels and output channels.

• reduction (int) – The reduction factor. Default: 16.

• activation (Optional[str]) – The type of activation function to use. Defaults to relu. Choices are relu, leaky_relu and elu.

forward(x)

The forward function of the squeeze and excitation layer.

Return type:

Tensor

training: bool

class yamle.models.operations.LSTM(input_size, hidden_size)

Bases: Module

This class is used to create a simple LSTM cell.

Parameters:

• input_size (int) – The number of input features.

• hidden_size (int) – The number of hidden features.

forward(x, h=None, c=None)

The forward function of the LSTM cell.

Parameters:

• x (torch.Tensor) – The input tensor of shape (batch_size, T, input_size).

• h (torch.Tensor, optional) – The hidden state of shape (batch_size, hidden_size).

• c (torch.Tensor, optional) – The cell state of shape (batch_size, hidden_size).

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

Tuple[torch.Tensor, torch.Tensor] – The complete hidden state of shape (batch_size, T, hidden_size) and the last hidden and cell state of shape (batch_size, hidden_size).

training: bool

class yamle.models.operations.Lambda(fn)

Bases: Module

This class is used to create a lambda layer.

Parameters:

fn (Callable) – The function to apply.

forward(x)

The forward function of the lambda layer.

Return type:

Tensor

training: bool

class yamle.models.operations.Reduction(dim, reduction='sum', alignment=False)

Bases: Module

Given a dimension this module implements a reduction operation.

Can choose between sum, mean, max, min and cat. If cat is chosen, the dimension size will be multiplied by the number of tensors.

Parameters:

• dim (int) – The dimension to reduce.

• reduction (str) – The reduction operation. Choices are: sum, mean, max, min and cat.

• alignment (bool) – If True, an alighment score will be computed for each tensor followed by a softmax.

forward(x)

The forward function of the reduction layer.

Return type:

Tensor

extra_repr()

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

training: bool

class yamle.models.operations.ScalarMultiplier(shape)

Bases: Module

This operation element-wise multiplies a tensor with learnable parameters.

The parameters are initialized to be 1. The operation can be enabled or disabled.

Parameters:

shape (tuple) – The shape of the learnable parameters.

forward(x)

The forward function of the scalar multiplier.

Parameters:

x (torch.Tensor) – The input tensor.

Return type:

Tensor

Returns:

torch.Tensor – The output tensor.

training: bool

enable()

Enable the scalar multiplier.

disable()

Disable the scalar multiplier.

extra_repr()

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str

class yamle.models.operations.ScalarAdder(shape)

Bases: ScalarMultiplier

This operation element-wise adds a tensor with learnable parameters.

The parameters are initialized to be 0. The operation can be enabled or disabled.

Parameters:

shape (tuple) – The shape of the learnable parameters.

training: bool

forward(x)

The forward function of the scalar adder.

Parameters:

x (torch.Tensor) – The input tensor.

Return type:

Tensor

Returns:

torch.Tensor – The output tensor.

class yamle.models.operations.LinearExtractor(inputs_dim, expansion_factor=None, hidden_dim=None, outputs_dim=1, depth=1, activation='ReLU', norm=False, normalization='batch', end_activation=False, end_normalization=False, residual=False)

Bases: Module

This module implements the linear neural network model.

It is used to be either input multiplexer or output demultiplexer. It takes in inputs_dim, expansion_factor and outputs_dim and produces a linear neural network model with activation and optionally norm layers in between. The number of layers is determined by depth.

Parameters:

• inputs_dim (int) – The input dimension for the input number of features.

• expansion_factor (Optional[float]) – The expansion factor for the hidden representation.

• hidden_dim (Optional[int]) – The hidden dimension for the hidden representation. If hidden_dim is None, it will be set to expansion_factor * inputs_dim.

• outputs_dim (int) – The output dimension for the output features.

• depth (int) – The number of layers in the sequence.

• activation (str) – The activation function to be used in the model. Default is ReLU.

• norm (bool) – Whether to use normalization layers in the model. Default is False.

• normalization (str) – The normalization layer to be used in the model. Default is BatchNorm1d.

• end_activation (bool) – Whether to use activation function at the end of the model. Default is False.

• end_normalization (bool) – Whether to use normalization layer at the end of the model. Default is False.

• residual (bool) – Whether to use residual connections in the model. Default is False.

training: bool

forward(x)

The forward function of the linear neural network model.

Return type:

Tensor

extra_repr()

The extra representation of the linear neural network model.

Return type:

str

class yamle.models.operations.Conv2dExtractor(input_channels, expansion_factor=None, hidden_dim=None, output_channels=1, depth=1, activation='ReLU', norm=False, normalization='batch', pool=False, pooling='max', end_activation=False, end_normalization=False, end_pooling=False, convolution='conv2d', residual=False, se=False)

Bases: Module

This module implements the convolutional neural network model.

It is used to be either input multiplexer or output demultiplexer. It takes in input_channels, expansion_factor and output_channels and produces a convolutional neural network model with activation and optionally norm layers in between. The number of layers is determined by depth. Global average pooling is optionally applied to the output through end_pooling.

Parameters:

• input_channels (int) – The input dimension for the input number of features.

• expansion_factor (Optional[float]) – The expansion factor for the hidden representation.

• hidden_dim (Optional[int]) – The hidden dimension for the hidden representation. If hidden_dim is None, it will be set to expansion_factor * inputs_dim.

• output_channels (int) – The output dimension for the number of output features.

• depth (int) – The number of layers in the sequence.

• activation (str) – The activation function to be used in the model. Default is ReLU.

• norm (bool) – Whether to use normalization layers in the model. Default is False.

• normalization (str) – The normalization layer to be used in the model. Default is BatchNorm2d.

• pool (bool) – Whether to use pooling layers in the model. Default is False.

• pooling (str) – The pooling layer to be used in the model. Default is max.

• end_activation (bool) – Whether to use activation function at the end of the model. Default is False.

• end_normalization (bool) – Whether to use normalization layer at the end of the model. Default is False.

• end_pooling (bool) – Whether to apply global average pooling to the output. Default is False.

• convolution (str) – The convolutional layer to be used. Default is conv2d. Choices are completely_separable, depthwise_separable, conv2d.

• residual (bool) – Whether to use residual connections in the model. Default is False.

• se (bool) – Whether to use squeeze and excitation in the model. Default is False.

training: bool

forward(x)

The forward function of the convolutional neural network model.

Return type:

Tensor

extra_repr()

The extra representation of the convolutional neural network model.

Return type:

str