torchwrench.nn package¶

class torchwrench.nn.Abs(*args: Any, **kwargs: Any)[source]¶

Bases: Module

Module version of abs().

forward(x: Tensor) → Tensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchwrench.nn.AdaptiveAvgPool1d(output_size: int | None | tuple[int | None, ...])[source]¶

Bases: _AdaptiveAvgPoolNd

Applies a 1D adaptive average pooling over an input signal composed of several input planes.

The output size is \(L_{out}\), for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size \(L_{out}\).

Shape:

• Input: \((N, C, L_{in})\) or \((C, L_{in})\).

• Output: \((N, C, L_{out})\) or \((C, L_{out})\), where \(L_{out}=\text{output\_size}\).

Examples:

>>> # target output size of 5
>>> m = nn.AdaptiveAvgPool1d(5)
>>> input = torch.randn(1, 64, 8)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

output_size : int | tuple[int]¶

class torchwrench.nn.AdaptiveAvgPool2d(output_size: int | None | tuple[int | None, ...])[source]¶

Bases: _AdaptiveAvgPoolNd

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

The output is of size H x W, for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size of the image of the form H x W.

Can be a tuple (H, W) or a single H for a square image H x H. H and W can be either a int, or None which means the size will be the same as that of the input.

Shape:

• Input: \((N, C, H_{in}, W_{in})\) or \((C, H_{in}, W_{in})\).

• Output: \((N, C, S_{0}, S_{1})\) or \((C, S_{0}, S_{1})\), where \(S=\text{output\_size}\).

Examples:

>>> # target output size of 5x7
>>> m = nn.AdaptiveAvgPool2d((5, 7))
>>> input = torch.randn(1, 64, 8, 9)
>>> output = m(input)
>>> # target output size of 7x7 (square)
>>> m = nn.AdaptiveAvgPool2d(7)
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)
>>> # target output size of 10x7
>>> m = nn.AdaptiveAvgPool2d((None, 7))
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

output_size : int | None | tuple[int | None, int | None]¶

class torchwrench.nn.AdaptiveAvgPool3d(output_size: int | None | tuple[int | None, ...])[source]¶

Bases: _AdaptiveAvgPoolNd

Applies a 3D adaptive average pooling over an input signal composed of several input planes.

The output is of size D x H x W, for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size of the form D x H x W.

Can be a tuple (D, H, W) or a single number D for a cube D x D x D. D, H and W can be either a int, or None which means the size will be the same as that of the input.

Shape:

• Input: \((N, C, D_{in}, H_{in}, W_{in})\) or \((C, D_{in}, H_{in}, W_{in})\).

• Output: \((N, C, S_{0}, S_{1}, S_{2})\) or \((C, S_{0}, S_{1}, S_{2})\), where \(S=\text{output\_size}\).

Examples:

>>> # target output size of 5x7x9
>>> m = nn.AdaptiveAvgPool3d((5, 7, 9))
>>> input = torch.randn(1, 64, 8, 9, 10)
>>> output = m(input)
>>> # target output size of 7x7x7 (cube)
>>> m = nn.AdaptiveAvgPool3d(7)
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)
>>> # target output size of 7x9x8
>>> m = nn.AdaptiveAvgPool3d((7, None, None))
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

output_size : int | None | tuple[int | None, int | None, int | None]¶

class torchwrench.nn.AdaptiveLogSoftmaxWithLoss(in_features: int, n_classes: int, cutoffs: Sequence[int], div_value: float = 4.0, head_bias: bool = False, device=None, dtype=None)[source]¶

Bases: Module

Efficient softmax approximation.

As described in Efficient softmax approximation for GPUs by Edouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, and Hervé Jégou.

Adaptive softmax is an approximate strategy for training models with large output spaces. It is most effective when the label distribution is highly imbalanced, for example in natural language modelling, where the word frequency distribution approximately follows the Zipf’s law.

Adaptive softmax partitions the labels into several clusters, according to their frequency. These clusters may contain different number of targets each. Additionally, clusters containing less frequent labels assign lower dimensional embeddings to those labels, which speeds up the computation. For each minibatch, only clusters for which at least one target is present are evaluated.

The idea is that the clusters which are accessed frequently (like the first one, containing most frequent labels), should also be cheap to compute – that is, contain a small number of assigned labels.

We highly recommend taking a look at the original paper for more details.

• cutoffs should be an ordered Sequence of integers sorted in the increasing order. It controls number of clusters and the partitioning of targets into clusters. For example setting cutoffs = [10, 100, 1000] means that first 10 targets will be assigned to the ‘head’ of the adaptive softmax, targets 11, 12, …, 100 will be assigned to the first cluster, and targets 101, 102, …, 1000 will be assigned to the second cluster, while targets 1001, 1002, …, n_classes - 1 will be assigned to the last, third cluster.

• div_value is used to compute the size of each additional cluster, which is given as \(\left\lfloor\frac{\texttt{in\_features}}{\texttt{div\_value}^{idx}}\right\rfloor\), where \(idx\) is the cluster index (with clusters for less frequent words having larger indices, and indices starting from \(1\)).

• head_bias if set to True, adds a bias term to the ‘head’ of the adaptive softmax. See paper for details. Set to False in the official implementation.

Warning

Labels passed as inputs to this module should be sorted according to their frequency. This means that the most frequent label should be represented by the index 0, and the least frequent label should be represented by the index n_classes - 1.

Note

This module returns a NamedTuple with output and loss fields. See further documentation for details.

Note

To compute log-probabilities for all classes, the log_prob method can be used.

Args:

in_features (int): Number of features in the input tensor n_classes (int): Number of classes in the dataset cutoffs (Sequence): Cutoffs used to assign targets to their buckets div_value (float, optional): value used as an exponent to compute sizes

of the clusters. Default: 4.0

head_bias (bool, optional): If True, adds a bias term to the ‘head’ of the

adaptive softmax. Default: False

Returns:

NamedTuple with output and loss fields:

• output is a Tensor of size N containing computed target log probabilities for each example

• loss is a Scalar representing the computed negative log likelihood loss

Shape:

• input: \((N, \texttt{in\_features})\) or \((\texttt{in\_features})\)

• target: \((N)\) or \(()\) where each value satisfies \(0 <= \texttt{target[i]} <= \texttt{n\_classes}\)

• output1: \((N)\) or \(()\)

• output2: Scalar

cutoffs : list[int]¶

div_value : float¶

forward(input_: Tensor, target_: Tensor) → _ASMoutput[source]¶

Runs the forward pass.

head : Linear¶

head_bias : bool¶

in_features : int¶

log_prob(input: Tensor) → Tensor[source]¶

Compute log probabilities for all \(\texttt{n\_classes}\).

Args:

input (Tensor): a minibatch of examples

Returns:

log-probabilities of for each class \(c\) in range \(0 <= c <= \texttt{n\_classes}\), where \(\texttt{n\_classes}\) is a parameter passed to AdaptiveLogSoftmaxWithLoss constructor.

Shape:

• Input: \((N, \texttt{in\_features})\)

• Output: \((N, \texttt{n\_classes})\)

n_classes : int¶

predict(input: Tensor) → Tensor[source]¶

Return the class with the highest probability for each example in the input minibatch.

This is equivalent to self.log_prob(input).argmax(dim=1), but is more efficient in some cases.

Args:

input (Tensor): a minibatch of examples

Returns:

output (Tensor): a class with the highest probability for each example

Shape:

• Input: \((N, \texttt{in\_features})\)

• Output: \((N)\)

reset_parameters() → None[source]¶

Resets parameters based on their initialization used in __init__.

tail : ModuleList¶

class torchwrench.nn.AdaptiveMaxPool1d(output_size: int | None | tuple[int | None, ...], return_indices: bool = False)[source]¶

Bases: _AdaptiveMaxPoolNd

Applies a 1D adaptive max pooling over an input signal composed of several input planes.

The output size is \(L_{out}\), for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size \(L_{out}\). return_indices: if True, will return the indices along with the outputs.

Useful to pass to nn.MaxUnpool1d. Default: False

Shape:

• Input: \((N, C, L_{in})\) or \((C, L_{in})\).

• Output: \((N, C, L_{out})\) or \((C, L_{out})\), where \(L_{out}=\text{output\_size}\).

Examples:

>>> # target output size of 5
>>> m = nn.AdaptiveMaxPool1d(5)
>>> input = torch.randn(1, 64, 8)
>>> output = m(input)

forward(input: Tensor)[source]¶

Runs the forward pass.

output_size : int | tuple[int]¶

class torchwrench.nn.AdaptiveMaxPool2d(output_size: int | None | tuple[int | None, ...], return_indices: bool = False)[source]¶

Bases: _AdaptiveMaxPoolNd

Applies a 2D adaptive max pooling over an input signal composed of several input planes.

The output is of size \(H_{out} \times W_{out}\), for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size of the image of the form \(H_{out} \times W_{out}\).

Can be a tuple \((H_{out}, W_{out})\) or a single \(H_{out}\) for a square image \(H_{out} \times H_{out}\). \(H_{out}\) and \(W_{out}\) can be either a int, or None which means the size will be the same as that of the input.

return_indices: if True, will return the indices along with the outputs.

Useful to pass to nn.MaxUnpool2d. Default: False

Shape:

• Input: \((N, C, H_{in}, W_{in})\) or \((C, H_{in}, W_{in})\).

• Output: \((N, C, H_{out}, W_{out})\) or \((C, H_{out}, W_{out})\), where \((H_{out}, W_{out})=\text{output\_size}\).

Examples:

>>> # target output size of 5x7
>>> m = nn.AdaptiveMaxPool2d((5, 7))
>>> input = torch.randn(1, 64, 8, 9)
>>> output = m(input)
>>> # target output size of 7x7 (square)
>>> m = nn.AdaptiveMaxPool2d(7)
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)
>>> # target output size of 10x7
>>> m = nn.AdaptiveMaxPool2d((None, 7))
>>> input = torch.randn(1, 64, 10, 9)
>>> output = m(input)

forward(input: Tensor)[source]¶

Runs the forward pass.

output_size : int | None | tuple[int | None, int | None]¶

class torchwrench.nn.AdaptiveMaxPool3d(output_size: int | None | tuple[int | None, ...], return_indices: bool = False)[source]¶

Bases: _AdaptiveMaxPoolNd

Applies a 3D adaptive max pooling over an input signal composed of several input planes.

The output is of size \(D_{out} \times H_{out} \times W_{out}\), for any input size. The number of output features is equal to the number of input planes.

Args:

output_size: the target output size of the image of the form \(D_{out} \times H_{out} \times W_{out}\).

Can be a tuple \((D_{out}, H_{out}, W_{out})\) or a single \(D_{out}\) for a cube \(D_{out} \times D_{out} \times D_{out}\). \(D_{out}\), \(H_{out}\) and \(W_{out}\) can be either a int, or None which means the size will be the same as that of the input.

return_indices: if True, will return the indices along with the outputs.

Useful to pass to nn.MaxUnpool3d. Default: False

Shape:

• Input: \((N, C, D_{in}, H_{in}, W_{in})\) or \((C, D_{in}, H_{in}, W_{in})\).

• Output: \((N, C, D_{out}, H_{out}, W_{out})\) or \((C, D_{out}, H_{out}, W_{out})\), where \((D_{out}, H_{out}, W_{out})=\text{output\_size}\).

Examples:

>>> # target output size of 5x7x9
>>> m = nn.AdaptiveMaxPool3d((5, 7, 9))
>>> input = torch.randn(1, 64, 8, 9, 10)
>>> output = m(input)
>>> # target output size of 7x7x7 (cube)
>>> m = nn.AdaptiveMaxPool3d(7)
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)
>>> # target output size of 7x9x8
>>> m = nn.AdaptiveMaxPool3d((7, None, None))
>>> input = torch.randn(1, 64, 10, 9, 8)
>>> output = m(input)

forward(input: Tensor)[source]¶

Runs the forward pass.

output_size : int | None | tuple[int | None, int | None, int | None]¶

class torchwrench.nn.AlphaDropout(p: float = 0.5, inplace: bool = False)[source]¶

Bases: _DropoutNd

Applies Alpha Dropout over the input.

Alpha Dropout is a type of Dropout that maintains the self-normalizing property. For an input with zero mean and unit standard deviation, the output of Alpha Dropout maintains the original mean and standard deviation of the input. Alpha Dropout goes hand-in-hand with SELU activation function, which ensures that the outputs have zero mean and unit standard deviation.

During training, it randomly masks some of the elements of the input tensor with probability p using samples from a bernoulli distribution. The elements to masked are randomized on every forward call, and scaled and shifted to maintain zero mean and unit standard deviation.

During evaluation the module simply computes an identity function.

More details can be found in the paper Self-Normalizing Neural Networks .

Args:

p (float): probability of an element to be dropped. Default: 0.5 inplace (bool, optional): If set to True, will do this operation

in-place

Shape:

• Input: \((*)\). Input can be of any shape

• Output: \((*)\). Output is of the same shape as input

Examples:

>>> m = nn.AlphaDropout(p=0.2)
>>> input = torch.randn(20, 16)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

class torchwrench.nn.Angle(*args: Any, **kwargs: Any)[source]¶

Bases: Module

Module version of angle().

forward(x: Tensor) → Tensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchwrench.nn.AsTensor(*, device: device | None | 'default' | 'cuda_if_available' | str | int = None, dtype: dtype | None | 'default' | str | DTypeEnum = None)[source]¶

Bases: Module

Module version of as_tensor().

extra_repr() → str[source]¶

Return 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.

forward(x: Any) → Tensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchwrench.nn.AvgPool1d(kernel_size: int | tuple[int], stride: int | tuple[int] = None, padding: int | tuple[int] = 0, ceil_mode: bool = False, count_include_pad: bool = True)[source]¶

Bases: _AvgPoolNd

Applies a 1D average pooling over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size \((N, C, L)\), output \((N, C, L_{out})\) and kernel_size \(k\) can be precisely described as:

\[\text{out}(N_i, C_j, l) = \frac{1}{k} \sum_{m=0}^{k-1} \text{input}(N_i, C_j, \text{stride} \times l + m)\]

If padding is non-zero, then the input is implicitly zero-padded on both sides for padding number of points.

Note:

When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding or the input. Sliding windows that would start in the right padded region are ignored.

Note

pad should be at most half of effective kernel size.

The parameters kernel_size, stride, padding can each be an int or a one-element tuple.

Args:

kernel_size: the size of the window stride: the stride of the window. Default value is kernel_size padding: implicit zero padding to be added on both sides ceil_mode: when True, will use ceil instead of floor to compute the output shape count_include_pad: when True, will include the zero-padding in the averaging calculation

Shape:

• Input: \((N, C, L_{in})\) or \((C, L_{in})\).

• Output: \((N, C, L_{out})\) or \((C, L_{out})\), where

  \[L_{out} = \left\lfloor \frac{L_{in} + 2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor\]

  Per the note above, if ceil_mode is True and \((L_{out} - 1) \times \text{stride} \geq L_{in} + \text{padding}\), we skip the last window as it would start in the right padded region, resulting in \(L_{out}\) being reduced by one.

Examples:

>>> # pool with window of size=3, stride=2
>>> m = nn.AvgPool1d(3, stride=2)
>>> m(torch.tensor([[[1., 2, 3, 4, 5, 6, 7]]]))
tensor([[[2., 4., 6.]]])

ceil_mode : bool¶

count_include_pad : bool¶

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

kernel_size : int | tuple[int]¶

padding : int | tuple[int]¶

stride : int | tuple[int]¶

class torchwrench.nn.AvgPool2d(kernel_size: int | tuple[int, int], stride: int | tuple[int, int] | None = None, padding: int | tuple[int, int] = 0, ceil_mode: bool = False, count_include_pad: bool = True, divisor_override: int | None = None)[source]¶

Bases: _AvgPoolNd

Applies a 2D average pooling over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size \((N, C, H, W)\), output \((N, C, H_{out}, W_{out})\) and kernel_size \((kH, kW)\) can be precisely described as:

\[out(N_i, C_j, h, w) = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} input(N_i, C_j, stride[0] \times h + m, stride[1] \times w + n)\]

If padding is non-zero, then the input is implicitly zero-padded on both sides for padding number of points.

Note:

When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding or the input. Sliding windows that would start in the right padded region are ignored.

Note

pad should be at most half of effective kernel size.

The parameters kernel_size, stride, padding can either be:

• a single int or a single-element tuple – in which case the same value is used for the height and width dimension

• a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Args:

kernel_size: the size of the window stride: the stride of the window. Default value is kernel_size padding: implicit zero padding to be added on both sides ceil_mode: when True, will use ceil instead of floor to compute the output shape count_include_pad: when True, will include the zero-padding in the averaging calculation divisor_override: if specified, it will be used as divisor, otherwise size of the pooling region will be used.

Shape:

• Input: \((N, C, H_{in}, W_{in})\) or \((C, H_{in}, W_{in})\).

• Output: \((N, C, H_{out}, W_{out})\) or \((C, H_{out}, W_{out})\), where

  \[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor\]
  \[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor\]

  Per the note above, if ceil_mode is True and \((H_{out} - 1)\times \text{stride}[0]\geq H_{in} + \text{padding}[0]\), we skip the last window as it would start in the bottom padded region, resulting in \(H_{out}\) being reduced by one.

  The same applies for \(W_{out}\).

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.AvgPool2d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.AvgPool2d((3, 2), stride=(2, 1))
>>> input = torch.randn(20, 16, 50, 32)
>>> output = m(input)

ceil_mode : bool¶

count_include_pad : bool¶

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

kernel_size : int | tuple[int, int]¶

padding : int | tuple[int, int]¶

stride : int | tuple[int, int]¶

class torchwrench.nn.AvgPool3d(kernel_size: int | tuple[int, int, int], stride: int | tuple[int, int, int] | None = None, padding: int | tuple[int, int, int] = 0, ceil_mode: bool = False, count_include_pad: bool = True, divisor_override: int | None = None)[source]¶

Bases: _AvgPoolNd

Applies a 3D average pooling over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size \((N, C, D, H, W)\), output \((N, C, D_{out}, H_{out}, W_{out})\) and kernel_size \((kD, kH, kW)\) can be precisely described as:

\[\begin{split}\begin{aligned} \text{out}(N_i, C_j, d, h, w) ={} & \sum_{k=0}^{kD-1} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} \\ & \frac{\text{input}(N_i, C_j, \text{stride}[0] \times d + k, \text{stride}[1] \times h + m, \text{stride}[2] \times w + n)} {kD \times kH \times kW} \end{aligned}\end{split}\]

If padding is non-zero, then the input is implicitly zero-padded on all three sides for padding number of points.

Note:

When ceil_mode=True, sliding windows are allowed to go off-bounds if they start within the left padding or the input. Sliding windows that would start in the right padded region are ignored.

Note

pad should be at most half of effective kernel size.

The parameters kernel_size, stride can either be:

• a single int – in which case the same value is used for the depth, height and width dimension

• a tuple of three ints – in which case, the first int is used for the depth dimension, the second int for the height dimension and the third int for the width dimension

Args:

kernel_size: the size of the window stride: the stride of the window. Default value is kernel_size padding: implicit zero padding to be added on all three sides ceil_mode: when True, will use ceil instead of floor to compute the output shape count_include_pad: when True, will include the zero-padding in the averaging calculation divisor_override: if specified, it will be used as divisor, otherwise kernel_size will be used

Shape:

• Input: \((N, C, D_{in}, H_{in}, W_{in})\) or \((C, D_{in}, H_{in}, W_{in})\).

• Output: \((N, C, D_{out}, H_{out}, W_{out})\) or \((C, D_{out}, H_{out}, W_{out})\), where

  \[D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor\]
  \[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor\]
  \[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{kernel\_size}[2]}{\text{stride}[2]} + 1\right\rfloor\]

  Per the note above, if ceil_mode is True and \((D_{out} - 1)\times \text{stride}[0]\geq D_{in} + \text{padding}[0]\), we skip the last window as it would start in the padded region, resulting in \(D_{out}\) being reduced by one.

  The same applies for \(W_{out}\) and \(H_{out}\).

Examples:

>>> # pool of square window of size=3, stride=2
>>> m = nn.AvgPool3d(3, stride=2)
>>> # pool of non-square window
>>> m = nn.AvgPool3d((3, 2, 2), stride=(2, 1, 2))
>>> input = torch.randn(20, 16, 50, 44, 31)
>>> output = m(input)

ceil_mode : bool¶

count_include_pad : bool¶

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

kernel_size : int | tuple[int, int, int]¶

padding : int | tuple[int, int, int]¶

stride : int | tuple[int, int, int]¶

class torchwrench.nn.BCELoss(weight: Tensor | None = None, size_average=None, reduce=None, reduction: str = 'mean')[source]¶

Bases: _WeightedLoss

Creates a criterion that measures the Binary Cross Entropy between the target and the input probabilities:

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_n \left[ y_n \cdot \log x_n + (1 - y_n) \cdot \log (1 - x_n) \right],\]

where \(N\) is the batch size. If reduction is not 'none' (default 'mean'), then

\[\begin{split}\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} \end{cases}\end{split}\]

This is used for measuring the error of a reconstruction in for example an auto-encoder. Note that the targets \(y\) should be numbers between 0 and 1.

Notice that if \(x_n\) is either 0 or 1, one of the log terms would be mathematically undefined in the above loss equation. PyTorch chooses to set \(\log (0) = -\infty\), since \(\lim_{x\to 0} \log (x) = -\infty\). However, an infinite term in the loss equation is not desirable for several reasons.

For one, if either \(y_n = 0\) or \((1 - y_n) = 0\), then we would be multiplying 0 with infinity. Secondly, if we have an infinite loss value, then we would also have an infinite term in our gradient, since \(\lim_{x\to 0} \frac{d}{dx} \log (x) = \infty\). This would make BCELoss’s backward method nonlinear with respect to \(x_n\), and using it for things like linear regression would not be straight-forward.

Our solution is that BCELoss clamps its log function outputs to be greater than or equal to -100. This way, we can always have a finite loss value and a linear backward method.

Args:

weight (Tensor, optional): a manual rescaling weight given to the loss

of each batch element. If given, has to be a Tensor of size nbatch.

size_average (bool, optional): Deprecated (see reduction). By default,

the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True

reduce (bool, optional): Deprecated (see reduction). By default, the

losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True

reduction (str, optional): Specifies the reduction to apply to the output:

'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

Shape:

• Input: \((*)\), where \(*\) means any number of dimensions.

• Target: \((*)\), same shape as the input.

• Output: scalar. If reduction is 'none', then \((*)\), same shape as input.

Examples:

>>> m = nn.Sigmoid()
>>> loss = nn.BCELoss()
>>> input = torch.randn(3, 2, requires_grad=True)
>>> target = torch.rand(3, 2, requires_grad=False)
>>> output = loss(m(input), target)
>>> output.backward()

forward(input: Tensor, target: Tensor) → Tensor[source]¶

Runs the forward pass.

class torchwrench.nn.BCEWithLogitsLoss(weight: Tensor | None = None, size_average=None, reduce=None, reduction: str = 'mean', pos_weight: Tensor | None = None)[source]¶

Bases: _Loss

This loss combines a Sigmoid layer and the BCELoss in one single class. This version is more numerically stable than using a plain Sigmoid followed by a BCELoss as, by combining the operations into one layer, we take advantage of the log-sum-exp trick for numerical stability.

The unreduced (i.e. with reduction set to 'none') loss can be described as:

\[\ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad l_n = - w_n \left[ y_n \cdot \log \sigma(x_n) + (1 - y_n) \cdot \log (1 - \sigma(x_n)) \right],\]

where \(N\) is the batch size. If reduction is not 'none' (default 'mean'), then

\[\begin{split}\ell(x, y) = \begin{cases} \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} \end{cases}\end{split}\]

This is used for measuring the error of a reconstruction in for example an auto-encoder. Note that the targets t[i] should be numbers between 0 and 1.

It’s possible to trade off recall and precision by adding weights to positive examples. In the case of multi-label classification the loss can be described as:

\[\ell_c(x, y) = L_c = \{l_{1,c},\dots,l_{N,c}\}^\top, \quad l_{n,c} = - w_{n,c} \left[ p_c y_{n,c} \cdot \log \sigma(x_{n,c}) + (1 - y_{n,c}) \cdot \log (1 - \sigma(x_{n,c})) \right],\]

where \(c\) is the class number (\(c > 1\) for multi-label binary classification, \(c = 1\) for single-label binary classification), \(n\) is the number of the sample in the batch and \(p_c\) is the weight of the positive answer for the class \(c\).

\(p_c > 1\) increases the recall, \(p_c < 1\) increases the precision.

For example, if a dataset contains 100 positive and 300 negative examples of a single class, then pos_weight for the class should be equal to \(\frac{300}{100}=3\). The loss would act as if the dataset contains \(3\times 100=300\) positive examples.

Examples:

>>> target = torch.ones([10, 64], dtype=torch.float32)  # 64 classes, batch size = 10
>>> output = torch.full([10, 64], 1.5)  # A prediction (logit)
>>> pos_weight = torch.ones([64])  # All weights are equal to 1
>>> criterion = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
>>> criterion(output, target)  # -log(sigmoid(1.5))
tensor(0.20...)

In the above example, the pos_weight tensor’s elements correspond to the 64 distinct classes in a multi-label binary classification scenario. Each element in pos_weight is designed to adjust the loss function based on the imbalance between negative and positive samples for the respective class. This approach is useful in datasets with varying levels of class imbalance, ensuring that the loss calculation accurately accounts for the distribution in each class.

Args:

weight (Tensor, optional): a manual rescaling weight given to the loss

of each batch element. If given, has to be a Tensor of size nbatch.

size_average (bool, optional): Deprecated (see reduction). By default,

the losses are averaged over each loss element in the batch. Note that for some losses, there are multiple elements per sample. If the field size_average is set to False, the losses are instead summed for each minibatch. Ignored when reduce is False. Default: True

reduce (bool, optional): Deprecated (see reduction). By default, the

losses are averaged or summed over observations for each minibatch depending on size_average. When reduce is False, returns a loss per batch element instead and ignores size_average. Default: True

reduction (str, optional): Specifies the reduction to apply to the output:

'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the sum of the output will be divided by the number of elements in the output, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction. Default: 'mean'

pos_weight (Tensor, optional): a weight of positive examples to be broadcasted with target.

Must be a tensor with equal size along the class dimension to the number of classes. Pay close attention to PyTorch’s broadcasting semantics in order to achieve the desired operations. For a target of size [B, C, H, W] (where B is batch size) pos_weight of size [B, C, H, W] will apply different pos_weights to each element of the batch or [C, H, W] the same pos_weights across the batch. To apply the same positive weight along all spatial dimensions for a 2D multi-class target [C, H, W] use: [C, 1, 1]. Default: None

Shape:

• Input: \((*)\), where \(*\) means any number of dimensions.

• Target: \((*)\), same shape as the input.

• Output: scalar. If reduction is 'none', then \((*)\), same shape as input.

Examples:

>>> loss = nn.BCEWithLogitsLoss()
>>> input = torch.randn(3, requires_grad=True)
>>> target = torch.empty(3).random_(2)
>>> output = loss(input, target)
>>> output.backward()

forward(input: Tensor, target: Tensor) → Tensor[source]¶

Runs the forward pass.

class torchwrench.nn.BatchNorm1d(num_features: int, eps: float = 1e-05, momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True, device=None, dtype=None)[source]¶

Bases: _BatchNorm

Applies Batch Normalization over a 2D or 3D input.

Method described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

\[y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

The mean and standard-deviation are calculated per-dimension over the mini-batches and \(\gamma\) and \(\beta\) are learnable parameter vectors of size C (where C is the number of features or channels of the input). By default, the elements of \(\gamma\) are set to 1 and the elements of \(\beta\) are set to 0. At train time in the forward pass, the variance is calculated via the biased estimator, equivalent to torch.var(input, correction=0). However, the value stored in the moving average of the variance is calculated via the unbiased estimator, equivalent to torch.var(input, correction=1).

Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default momentum of 0.1.

If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.

Note

This momentum argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is \(\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t\), where \(\hat{x}\) is the estimated statistic and \(x_t\) is the new observed value.

Because the Batch Normalization is done over the C dimension, computing statistics on (N, L) slices, it’s common terminology to call this Temporal Batch Normalization.

Args:

num_features: number of features or channels \(C\) of the input eps: a value added to the denominator for numerical stability.

Default: 1e-5

momentum: the value used for the running_mean and running_var

computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1

affine: a boolean value that when set to True, this module has

learnable affine parameters. Default: True

track_running_stats: a boolean value that when set to True, this

module tracks the running mean and variance, and when set to False, this module does not track such statistics, and initializes statistics buffers running_mean and running_var as None. When these buffers are None, this module always uses batch statistics. in both training and eval modes. Default: True

Shape:

• Input: \((N, C)\) or \((N, C, L)\), where \(N\) is the batch size, \(C\) is the number of features or channels, and \(L\) is the sequence length

• Output: \((N, C)\) or \((N, C, L)\) (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm1d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm1d(100, affine=False)
>>> input = torch.randn(20, 100)
>>> output = m(input)

class torchwrench.nn.BatchNorm2d(num_features: int, eps: float = 1e-05, momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True, device=None, dtype=None)[source]¶

Bases: _BatchNorm

Applies Batch Normalization over a 4D input.

4D is a mini-batch of 2D inputs with additional channel dimension. Method described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

The mean and standard-deviation are calculated per-dimension over the mini-batches and \(\gamma\) and \(\beta\) are learnable parameter vectors of size C (where C is the input size). By default, the elements of \(\gamma\) are set to 1 and the elements of \(\beta\) are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent to torch.var(input, correction=0). However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent to torch.var(input, correction=1).

Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default momentum of 0.1.

If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.

Note

This momentum argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is \(\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t\), where \(\hat{x}\) is the estimated statistic and \(x_t\) is the new observed value.

Because the Batch Normalization is done over the C dimension, computing statistics on (N, H, W) slices, it’s common terminology to call this Spatial Batch Normalization.

Args:

num_features: \(C\) from an expected input of size

\((N, C, H, W)\)

eps: a value added to the denominator for numerical stability.

Default: 1e-5

momentum: the value used for the running_mean and running_var

computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1

affine: a boolean value that when set to True, this module has

learnable affine parameters. Default: True

track_running_stats: a boolean value that when set to True, this

module tracks the running mean and variance, and when set to False, this module does not track such statistics, and initializes statistics buffers running_mean and running_var as None. When these buffers are None, this module always uses batch statistics. in both training and eval modes. Default: True

Shape:

• Input: \((N, C, H, W)\)

• Output: \((N, C, H, W)\) (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm2d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm2d(100, affine=False)
>>> input = torch.randn(20, 100, 35, 45)
>>> output = m(input)

class torchwrench.nn.BatchNorm3d(num_features: int, eps: float = 1e-05, momentum: float | None = 0.1, affine: bool = True, track_running_stats: bool = True, device=None, dtype=None)[source]¶

Bases: _BatchNorm

Applies Batch Normalization over a 5D input.

5D is a mini-batch of 3D inputs with additional channel dimension as described in the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

\[y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta\]

The mean and standard-deviation are calculated per-dimension over the mini-batches and \(\gamma\) and \(\beta\) are learnable parameter vectors of size C (where C is the input size). By default, the elements of \(\gamma\) are set to 1 and the elements of \(\beta\) are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent to torch.var(input, correction=0). However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent to torch.var(input, correction=1).

Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default momentum of 0.1.

If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.

Note

This momentum argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here is \(\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t\), where \(\hat{x}\) is the estimated statistic and \(x_t\) is the new observed value.

Because the Batch Normalization is done over the C dimension, computing statistics on (N, D, H, W) slices, it’s common terminology to call this Volumetric Batch Normalization or Spatio-temporal Batch Normalization.

Args:

num_features: \(C\) from an expected input of size

\((N, C, D, H, W)\)

eps: a value added to the denominator for numerical stability.

Default: 1e-5

momentum: the value used for the running_mean and running_var

computation. Can be set to None for cumulative moving average (i.e. simple average). Default: 0.1

affine: a boolean value that when set to True, this module has

learnable affine parameters. Default: True

track_running_stats: a boolean value that when set to True, this

module tracks the running mean and variance, and when set to False, this module does not track such statistics, and initializes statistics buffers running_mean and running_var as None. When these buffers are None, this module always uses batch statistics. in both training and eval modes. Default: True

Shape:

• Input: \((N, C, D, H, W)\)

• Output: \((N, C, D, H, W)\) (same shape as input)

Examples:

>>> # With Learnable Parameters
>>> m = nn.BatchNorm3d(100)
>>> # Without Learnable Parameters
>>> m = nn.BatchNorm3d(100, affine=False)
>>> input = torch.randn(20, 100, 35, 45, 10)
>>> output = m(input)

class torchwrench.nn.Bilinear(in1_features: int, in2_features: int, out_features: int, bias: bool = True, device=None, dtype=None)[source]¶

Bases: Module

Applies a bilinear transformation to the incoming data: \(y = x_1^T A x_2 + b\).

Args:

in1_features: size of each first input sample, must be > 0 in2_features: size of each second input sample, must be > 0 out_features: size of each output sample, must be > 0 bias: If set to False, the layer will not learn an additive bias.

Default: True

Shape:

• Input1: \((*, H_\text{in1})\) where \(H_\text{in1}=\text{in1\_features}\) and \(*\) means any number of additional dimensions including none. All but the last dimension of the inputs should be the same.

• Input2: \((*, H_\text{in2})\) where \(H_\text{in2}=\text{in2\_features}\).

• Output: \((*, H_\text{out})\) where \(H_\text{out}=\text{out\_features}\) and all but the last dimension are the same shape as the input.

Attributes:

weight: the learnable weights of the module of shape

\((\text{out\_features}, \text{in1\_features}, \text{in2\_features})\). The values are initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\), where \(k = \frac{1}{\text{in1\_features}}\)

bias: the learnable bias of the module of shape \((\text{out\_features})\).

If bias is True, the values are initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\), where \(k = \frac{1}{\text{in1\_features}}\)

Examples:

>>> m = nn.Bilinear(20, 30, 40)
>>> input1 = torch.randn(128, 20)
>>> input2 = torch.randn(128, 30)
>>> output = m(input1, input2)
>>> print(output.size())
torch.Size([128, 40])

extra_repr() → str[source]¶

Return the extra representation of the module.

forward(input1: Tensor, input2: Tensor) → Tensor[source]¶

Runs the forward pass.

in1_features : int¶

in2_features : int¶

out_features : int¶

reset_parameters() → None[source]¶

Resets parameters based on their initialization used in __init__.

weight : Tensor¶

class torchwrench.nn.CELU(alpha: float = 1.0, inplace: bool = False)[source]¶

Bases: Module

Applies the CELU function element-wise.

\[\text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))\]

More details can be found in the paper Continuously Differentiable Exponential Linear Units .

Args:

alpha: the \(\alpha\) value for the CELU formulation. Default: 1.0 inplace: can optionally do the operation in-place. Default: False

Shape:

• Input: \((*)\), where \(*\) means any number of dimensions.

• Output: \((*)\), same shape as the input.

Examples:

>>> m = nn.CELU()
>>> input = torch.randn(2)
>>> output = m(input)

alpha : float¶

extra_repr() → str[source]¶

Return the extra representation of the module.

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

inplace : bool¶

class torchwrench.nn.CTCLoss(blank: int = 0, reduction: str = 'mean', zero_infinity: bool = False)[source]¶

Bases: _Loss

The Connectionist Temporal Classification loss.

Calculates loss between a continuous (unsegmented) time series and a target sequence. CTCLoss sums over the probability of possible alignments of input to target, producing a loss value which is differentiable with respect to each input node. The alignment of input to target is assumed to be “many-to-one”, which limits the length of the target sequence such that it must be \(\leq\) the input length.

Args:

blank (int, optional): blank label. Default \(0\). reduction (str, optional): Specifies the reduction to apply to the output:

'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the output losses will be divided by the target lengths and then the mean over the batch is taken, 'sum': the output losses will be summed. Default: 'mean'

zero_infinity (bool, optional):

Whether to zero infinite losses and the associated gradients. Default: False Infinite losses mainly occur when the inputs are too short to be aligned to the targets.

Shape:

• Log_probs: Tensor of size \((T, N, C)\) or \((T, C)\), where \(T = \text{input length}\), \(N = \text{batch size}\), and \(C = \text{number of classes (including blank)}\). The logarithmized probabilities of the outputs (e.g. obtained with torch.nn.functional.log_softmax()).

• Targets: Tensor of size \((N, S)\) or \((\operatorname{sum}(\text{target\_lengths}))\), where \(N = \text{batch size}\) and \(S = \text{max target length, if shape is } (N, S)\). It represents the target sequences. Each element in the target sequence is a class index. And the target index cannot be blank (default=0). In the \((N, S)\) form, targets are padded to the length of the longest sequence, and stacked. In the \((\operatorname{sum}(\text{target\_lengths}))\) form, the targets are assumed to be un-padded and concatenated within 1 dimension.

• Input_lengths: Tuple or tensor of size \((N)\) or \(()\), where \(N = \text{batch size}\). It represents the lengths of the inputs (must each be \(\leq T\)). And the lengths are specified for each sequence to achieve masking under the assumption that sequences are padded to equal lengths.

• Target_lengths: Tuple or tensor of size \((N)\) or \(()\), where \(N = \text{batch size}\). It represents lengths of the targets. Lengths are specified for each sequence to achieve masking under the assumption that sequences are padded to equal lengths. If target shape is \((N,S)\), target_lengths are effectively the stop index \(s_n\) for each target sequence, such that target_n = targets[n,0:s_n] for each target in a batch. Lengths must each be \(\leq S\) If the targets are given as a 1d tensor that is the concatenation of individual targets, the target_lengths must add up to the total length of the tensor.

• Output: scalar if reduction is 'mean' (default) or 'sum'. If reduction is 'none', then \((N)\) if input is batched or \(()\) if input is unbatched, where \(N = \text{batch size}\).

Examples:

>>> # Target are to be padded
>>> T = 50  # Input sequence length
>>> C = 20  # Number of classes (including blank)
>>> N = 16  # Batch size
>>> S = 30  # Target sequence length of longest target in batch (padding length)
>>> S_min = 10  # Minimum target length, for demonstration purposes
>>>
>>> # Initialize random batch of input vectors, for *size = (T,N,C)
>>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()
>>>
>>> # Initialize random batch of targets (0 = blank, 1:C = classes)
>>> target = torch.randint(low=1, high=C, size=(N, S), dtype=torch.long)
>>>
>>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)
>>> target_lengths = torch.randint(
...     low=S_min,
...     high=S,
...     size=(N,),
...     dtype=torch.long,
... )
>>> ctc_loss = nn.CTCLoss()
>>> loss = ctc_loss(input, target, input_lengths, target_lengths)
>>> loss.backward()
>>>
>>>
>>> # Target are to be un-padded
>>> T = 50  # Input sequence length
>>> C = 20  # Number of classes (including blank)
>>> N = 16  # Batch size
>>>
>>> # Initialize random batch of input vectors, for *size = (T,N,C)
>>> input = torch.randn(T, N, C).log_softmax(2).detach().requires_grad_()
>>> input_lengths = torch.full(size=(N,), fill_value=T, dtype=torch.long)
>>>
>>> # Initialize random batch of targets (0 = blank, 1:C = classes)
>>> target_lengths = torch.randint(low=1, high=T, size=(N,), dtype=torch.long)
>>> target = torch.randint(
...     low=1,
...     high=C,
...     size=(sum(target_lengths),),
...     dtype=torch.long,
... )
>>> ctc_loss = nn.CTCLoss()
>>> loss = ctc_loss(input, target, input_lengths, target_lengths)
>>> loss.backward()
>>>
>>>
>>> # Target are to be un-padded and unbatched (effectively N=1)
>>> T = 50  # Input sequence length
>>> C = 20  # Number of classes (including blank)
>>>
>>> # Initialize random batch of input vectors, for *size = (T,C)
>>> # xdoctest: +SKIP("FIXME: error in doctest")
>>> input = torch.randn(T, C).log_softmax(1).detach().requires_grad_()
>>> input_lengths = torch.tensor(T, dtype=torch.long)
>>>
>>> # Initialize random batch of targets (0 = blank, 1:C = classes)
>>> target_lengths = torch.randint(low=1, high=T, size=(), dtype=torch.long)
>>> target = torch.randint(
...     low=1,
...     high=C,
...     size=(target_lengths,),
...     dtype=torch.long,
... )
>>> ctc_loss = nn.CTCLoss()
>>> loss = ctc_loss(input, target, input_lengths, target_lengths)
>>> loss.backward()

Reference:

A. Graves et al.: Connectionist Temporal Classification: Labelling Unsegmented Sequence Data with Recurrent Neural Networks: https://www.cs.toronto.edu/~graves/icml_2006.pdf

Note:

In order to use CuDNN, the following must be satisfied: the targets must be in concatenated format, all input_lengths must be T. \(blank=0\), target_lengths \(\leq 256\), the integer arguments must be of dtype torch.int32, and the log_probs itself must be of dtype torch.float32.

The regular implementation uses the (more common in PyTorch) torch.long dtype.

Note:

In some circumstances when using the CUDA backend with CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. Please see the notes on /notes/randomness for background.

blank : int¶

forward(log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor) → Tensor[source]¶

Runs the forward pass.

zero_infinity : bool¶

class torchwrench.nn.ChannelShuffle(groups: int)[source]¶

Bases: Module

Divides and rearranges the channels in a tensor.

This operation divides the channels in a tensor of shape \((N, C, *)\) into g groups as \((N, \frac{C}{g}, g, *)\) and shuffles them, while retaining the original tensor shape in the final output.

Args:

groups (int): number of groups to divide channels in.

Examples:

>>> channel_shuffle = nn.ChannelShuffle(2)
>>> input = torch.arange(1, 17, dtype=torch.float32).view(1, 4, 2, 2)
>>> input
tensor([[[[ 1.,  2.],
          [ 3.,  4.]],
         [[ 5.,  6.],
          [ 7.,  8.]],
         [[ 9., 10.],
          [11., 12.]],
         [[13., 14.],
          [15., 16.]]]])
>>> output = channel_shuffle(input)
>>> output
tensor([[[[ 1.,  2.],
          [ 3.,  4.]],
         [[ 9., 10.],
          [11., 12.]],
         [[ 5.,  6.],
          [ 7.,  8.]],
         [[13., 14.],
          [15., 16.]]]])

extra_repr() → str[source]¶

Return the extra representation of the module.

forward(input: Tensor) → Tensor[source]¶

Runs the forward pass.

groups : int¶

class torchwrench.nn.ConstantPad1d(padding: int | tuple[int, int], value: float)[source]¶

Bases: _ConstantPadNd

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Args:

padding (int, tuple): the size of the padding. If is int, uses the same

padding in both boundaries. If a 2-tuple, uses (\(\text{padding\_left}\), \(\text{padding\_right}\))

Shape:

• Input: \((C, W_{in})\) or \((N, C, W_{in})\).

• Output: \((C, W_{out})\) or \((N, C, W_{out})\), where

  \(W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}\)

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> m = nn.ConstantPad1d(2, 3.5)
>>> input = torch.randn(1, 2, 4)
>>> input
tensor([[[-1.0491, -0.7152, -0.0749,  0.8530],
         [-1.3287,  1.8966,  0.1466, -0.2771]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000, -1.0491, -0.7152, -0.0749,  0.8530,  3.5000,
           3.5000],
         [ 3.5000,  3.5000, -1.3287,  1.8966,  0.1466, -0.2771,  3.5000,
           3.5000]]])
>>> m = nn.ConstantPad1d(2, 3.5)
>>> input = torch.randn(1, 2, 3)
>>> input
tensor([[[ 1.6616,  1.4523, -1.1255],
         [-3.6372,  0.1182, -1.8652]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000,  3.5000],
         [ 3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000,  3.5000]]])
>>> # using different paddings for different sides
>>> m = nn.ConstantPad1d((3, 1), 3.5)
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  1.6616,  1.4523, -1.1255,  3.5000],
         [ 3.5000,  3.5000,  3.5000, -3.6372,  0.1182, -1.8652,  3.5000]]])

padding : tuple[int, int]¶

class torchwrench.nn.ConstantPad2d(padding: int | tuple[int, int, int, int], value: float)[source]¶

Bases: _ConstantPadNd

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Args:

padding (int, tuple): the size of the padding. If is int, uses the same

padding in all boundaries. If a 4-tuple, uses (\(\text{padding\_left}\), \(\text{padding\_right}\), \(\text{padding\_top}\), \(\text{padding\_bottom}\))

Shape:

• Input: \((N, C, H_{in}, W_{in})\) or \((C, H_{in}, W_{in})\).

• Output: \((N, C, H_{out}, W_{out})\) or \((C, H_{out}, W_{out})\), where

  \(H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}\)

  \(W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}\)

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> m = nn.ConstantPad2d(2, 3.5)
>>> input = torch.randn(1, 2, 2)
>>> input
tensor([[[ 1.6585,  0.4320],
         [-0.8701, -0.4649]]])
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  1.6585,  0.4320,  3.5000,  3.5000],
         [ 3.5000,  3.5000, -0.8701, -0.4649,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])
>>> # using different paddings for different sides
>>> m = nn.ConstantPad2d((3, 0, 2, 1), 3.5)
>>> m(input)
tensor([[[ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000],
         [ 3.5000,  3.5000,  3.5000,  1.6585,  0.4320],
         [ 3.5000,  3.5000,  3.5000, -0.8701, -0.4649],
         [ 3.5000,  3.5000,  3.5000,  3.5000,  3.5000]]])

padding : tuple[int, int, int, int]¶

class torchwrench.nn.ConstantPad3d(padding: int | tuple[int, int, int, int, int, int], value: float)[source]¶

Bases: _ConstantPadNd

Pads the input tensor boundaries with a constant value.

For N-dimensional padding, use torch.nn.functional.pad().

Args:

padding (int, tuple): the size of the padding. If is int, uses the same

padding in all boundaries. If a 6-tuple, uses (\(\text{padding\_left}\), \(\text{padding\_right}\), \(\text{padding\_top}\), \(\text{padding\_bottom}\), \(\text{padding\_front}\), \(\text{padding\_back}\))

Shape:

• Input: \((N, C, D_{in}, H_{in}, W_{in})\) or \((C, D_{in}, H_{in}, W_{in})\).

• Output: \((N, C, D_{out}, H_{out}, W_{out})\) or \((C, D_{out}, H_{out}, W_{out})\), where

  \(D_{out} = D_{in} + \text{padding\_front} + \text{padding\_back}\)

  \(H_{out} = H_{in} + \text{padding\_top} + \text{padding\_bottom}\)

  \(W_{out} = W_{in} + \text{padding\_left} + \text{padding\_right}\)

Examples:

>>> m = nn.ConstantPad3d(3, 3.5)
>>> input = torch.randn(16, 3, 10, 20, 30)
>>> output = m(input)
>>> # using different paddings for different sides
>>> m = nn.ConstantPad3d((3, 3, 6, 6, 0, 1), 3.5)
>>> output = m(input)

padding : tuple[int, int, int, int, int, int]¶

class torchwrench.nn.Conv1d(in_channels: int, out_channels: int, kernel_size: int | tuple[int], stride: int | tuple[int] = 1, padding: str | int | tuple[int] = 0, dilation: int | tuple[int] = 1, groups: int = 1, bias: bool = True, padding_mode: 'zeros' | 'reflect' | 'replicate' | 'circular' = 'zeros', device=None, dtype=None)[source]¶

Bases: _ConvNd

Applies a 1D convolution over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size \((N, C_{\text{in}}, L)\) and output \((N, C_{\text{out}}, L_{\text{out}})\) can be precisely described as:

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{in} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)\]

where \(\star\) is the valid cross-correlation operator, \(N\) is a batch size, \(C\) denotes a number of channels, \(L\) is a length of signal sequence.

This module supports TensorFloat32.

On certain ROCm devices, when using float16 inputs this module will use different precision for backward.

• stride controls the stride for the cross-correlation, a single number or a one-element tuple.

• padding controls the amount of padding applied to the input. It can be either a string {‘valid’, ‘same’} or a tuple of ints giving the amount of implicit padding applied on both sides.

• dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.

• groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,

  • At groups=1, all inputs are convolved to all outputs.

  • At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels and producing half the output channels, and both subsequently concatenated.

  • At groups= in_channels, each input channel is convolved with its own set of filters (of size \(\frac{\text{out\_channels}}{\text{in\_channels}}\)).

Note:

When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, this operation is also known as a “depthwise convolution”.

In other words, for an input of size \((N, C_{in}, L_{in})\), a depthwise convolution with a depthwise multiplier K can be performed with the arguments \((C_\text{in}=C_\text{in}, C_\text{out}=C_\text{in} \times \text{K}, ..., \text{groups}=C_\text{in})\).

Note:

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Note:

padding='valid' is the same as no padding. padding='same' pads the input so the output has the shape as the input. However, this mode doesn’t support any stride values other than 1.

Note:

This module supports complex data types i.e. complex32, complex64, complex128.

Args:

in_channels (int): Number of channels in the input image out_channels (int): Number of channels produced by the convolution kernel_size (int or tuple): Size of the convolving kernel stride (int or tuple, optional): Stride of the convolution. Default: 1 padding (int, tuple or str, optional): Padding added to both sides of

the input. Default: 0

dilation (int or tuple, optional): Spacing between kernel

elements. Default: 1

groups (int, optional): Number of blocked connections from input

channels to output channels. Default: 1

bias (bool, optional): If True, adds a learnable bias to the

output. Default: True

padding_mode (str, optional): 'zeros', 'reflect',

'replicate' or 'circular'. Default: 'zeros'

Shape:

• Input: \((N, C_{in}, L_{in})\) or \((C_{in}, L_{in})\)

• Output: \((N, C_{out}, L_{out})\) or \((C_{out}, L_{out})\), where

  \[L_{out} = \left\lfloor\frac{L_{in} + 2 \times \text{padding} - \text{dilation} \times (\text{kernel\_size} - 1) - 1}{\text{stride}} + 1\right\rfloor\]

Attributes:

weight (Tensor): the learnable weights of the module of shape

\((\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}}, \text{kernel\_size})\). The values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{groups}{C_\text{in} * \text{kernel\_size}}\)

bias (Tensor): the learnable bias of the module of shape

(out_channels). If bias is True, then the values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{groups}{C_\text{in} * \text{kernel\_size}}\)

Examples:

>>> m = nn.Conv1d(16, 33, 3, stride=2)
>>> input = torch.randn(20, 16, 50)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchwrench.nn.Conv2d(in_channels: int, out_channels: int, kernel_size: int | tuple[int, int], stride: int | tuple[int, int] = 1, padding: str | int | tuple[int, int] = 0, dilation: int | tuple[int, int] = 1, groups: int = 1, bias: bool = True, padding_mode: 'zeros' | 'reflect' | 'replicate' | 'circular' = 'zeros', device=None, dtype=None)[source]¶

Bases: _ConvNd

Applies a 2D convolution over an input signal composed of several input planes.

In the simplest case, the output value of the layer with input size \((N, C_{\text{in}}, H, W)\) and output \((N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})\) can be precisely described as:

\[\text{out}(N_i, C_{\text{out}_j}) = \text{bias}(C_{\text{out}_j}) + \sum_{k = 0}^{C_{\text{in}} - 1} \text{weight}(C_{\text{out}_j}, k) \star \text{input}(N_i, k)\]

where \(\star\) is the valid 2D cross-correlation operator, \(N\) is a batch size, \(C\) denotes a number of channels, \(H\) is a height of input planes in pixels, and \(W\) is width in pixels.

This module supports TensorFloat32.

On certain ROCm devices, when using float16 inputs this module will use different precision for backward.

• stride controls the stride for the cross-correlation, a single number or a tuple.

• padding controls the amount of padding applied to the input. It can be either a string {‘valid’, ‘same’} or an int / a tuple of ints giving the amount of implicit padding applied on both sides.

• dilation controls the spacing between the kernel points; also known as the à trous algorithm. It is harder to describe, but this link has a nice visualization of what dilation does.

• groups controls the connections between inputs and outputs. in_channels and out_channels must both be divisible by groups. For example,

  • At groups=1, all inputs are convolved to all outputs.

  • At groups=2, the operation becomes equivalent to having two conv layers side by side, each seeing half the input channels and producing half the output channels, and both subsequently concatenated.

  • At groups= in_channels, each input channel is convolved with its own set of filters (of size \(\frac{\text{out\_channels}}{\text{in\_channels}}\)).

The parameters kernel_size, stride, padding, dilation can either be:

• a single int – in which case the same value is used for the height and width dimension

• a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension

Note:

When groups == in_channels and out_channels == K * in_channels, where K is a positive integer, this operation is also known as a “depthwise convolution”.

In other words, for an input of size \((N, C_{in}, L_{in})\), a depthwise convolution with a depthwise multiplier K can be performed with the arguments \((C_\text{in}=C_\text{in}, C_\text{out}=C_\text{in} \times \text{K}, ..., \text{groups}=C_\text{in})\).

Note:

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Note:

padding='valid' is the same as no padding. padding='same' pads the input so the output has the shape as the input. However, this mode doesn’t support any stride values other than 1.

Note:

This module supports complex data types i.e. complex32, complex64, complex128.

Args:

in_channels (int): Number of channels in the input image out_channels (int): Number of channels produced by the convolution kernel_size (int or tuple): Size of the convolving kernel stride (int or tuple, optional): Stride of the convolution. Default: 1 padding (int, tuple or str, optional): Padding added to all four sides of

the input. Default: 0

dilation (int or tuple, optional): Spacing between kernel elements. Default: 1 groups (int, optional): Number of blocked connections from input

channels to output channels. Default: 1

bias (bool, optional): If True, adds a learnable bias to the

output. Default: True

padding_mode (str, optional): 'zeros', 'reflect',

'replicate' or 'circular'. Default: 'zeros'

Shape:

• Input: \((N, C_{in}, H_{in}, W_{in})\) or \((C_{in}, H_{in}, W_{in})\)

• Output: \((N, C_{out}, H_{out}, W_{out})\) or \((C_{out}, H_{out}, W_{out})\), where

  \[H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor\]
  \[W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor\]

Attributes:

weight (Tensor): the learnable weights of the module of shape

\((\text{out\_channels}, \frac{\text{in\_channels}}{\text{groups}},\) \(\text{kernel\_size[0]}, \text{kernel\_size[1]})\). The values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}\)

bias (Tensor): the learnable bias of the module of shape

(out_channels). If bias is True, then the values of these weights are sampled from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where \(k = \frac{groups}{C_\text{in} * \prod_{i=0}^{1}\text{kernel\_size}[i]}\)

Examples:

>>> # With square kernels and equal stride
>>> m = nn.Conv2d(16, 33, 3, stride=2)
>>> # non-square kernels and unequal stride and with padding
>>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2))
>>> # non-square kernels and unequal stride and with padding and dilation
>>> m = nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))
>>> input = torch.randn(20, 16, 50, 100)
>>> output = m(input)

forward(input: Tensor) → Tensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.