holocron.nn

An addition to the torch.nn module of Pytorch to extend the range of neural networks building blocks.

Non-linear activations

HardMish

HardMish(inplace: bool = False)

Bases: _Activation

Implements the Hard Mish activation module from "H-Mish".

This activation is computed as follows:

\[ f(x) = \frac{x}{2} \cdot \min(2, \max(0, x + 2)) \]

PARAMETER	DESCRIPTION
inplace	should the operation be performed inplace TYPE: bool DEFAULT: False

Source code in holocron/nn/modules/activation.py

def __init__(self, inplace: bool = False) -> None:
    super().__init__()
    self.inplace: bool = inplace

NLReLU

NLReLU(beta: float = 1.0, inplace: bool = False)

Bases: _Activation

Implements the Natural-Logarithm ReLU activation module from "Natural-Logarithm-Rectified Activation Function in Convolutional Neural Networks".

This activation is computed as follows:

\[ f(x) = ln(1 + \beta \cdot max(0, x)) \]

PARAMETER	DESCRIPTION
beta	beta used for NReLU TYPE: float DEFAULT: 1.0
inplace	should the operation be performed inplace TYPE: bool DEFAULT: False

Source code in holocron/nn/modules/activation.py

def __init__(self, beta: float = 1.0, inplace: bool = False) -> None:
    super().__init__(inplace)
    self.beta: float = beta

FReLU

FReLU(in_channels: int, kernel_size: int = 3)

Bases: Module

Implements the Funnel activation module from "Funnel Activation for Visual Recognition".

This activation is computed as follows:

\[ f(x) = max(\mathbb{T}(x), x) \]

where the \(\mathbb{T}\) is the spatial contextual feature extraction. It is a convolution filter of size kernel_size, same padding and groups equal to the number of input channels, followed by a batch normalization.

PARAMETER	DESCRIPTION
in_channels	number of input channels TYPE: int
kernel_size	size of the convolution filter TYPE: int DEFAULT: 3

Source code in holocron/nn/modules/activation.py

def __init__(self, in_channels: int, kernel_size: int = 3) -> None:
    super().__init__()
    self.conv = nn.Conv2d(in_channels, in_channels, kernel_size, padding=kernel_size // 2, groups=in_channels)
    self.bn = nn.BatchNorm2d(in_channels)

Loss functions

Check each loss's input contract before using it. N is the batch size and K is the number of classes; ... represents spatial dimensions, which are optional for FocalLoss and PolyLoss.

Loss	Input	Target
FocalLoss	Raw logits, shape (N, K, ...)	Class indices, shape (N, ...), dtype torch.int64
PolyLoss	Raw logits, shape (N, K, ...)	Class indices with dtype torch.int64, or soft class probabilities with the same shape as the logits
DiceLoss	Class probabilities, shape (N, K, ...)	One-hot or soft targets with the same shape as the input

DiceLoss requires at least one spatial dimension and returns a scalar. It does not accept ignore_index or reduction arguments.

For hard targets, class indices must be in [0, K), except for ignored values supported by the selected loss. A segmentation target has shape (N, H, W), without a class channel. Casting a soft target to torch.int64 does not turn it into class indices.

FocalLoss(ignore_index=-100) supports ignored hard targets outside the class range and excludes them from the loss and gradient. PolyLoss currently does not: a hard target containing -100 or 255 outside [0, K) raises RuntimeError, even when it matches ignore_index. Its in-range ignore behavior and reduction="none" limits are documented below.

Loss

Loss(weight: float | list[float] | Tensor | None = None, ignore_index: int = -100, reduction: str = 'mean')

Bases: Module

Base loss class.

PARAMETER	DESCRIPTION
weight	class weight for loss computation TYPE: float | list[float] | Tensor | None DEFAULT: None
ignore_index	specifies target value that is ignored and do not contribute to gradient TYPE: int DEFAULT: -100
reduction	type of reduction to apply to the final loss TYPE: str DEFAULT: 'mean'

RAISES	DESCRIPTION
NotImplementedError	if the reduction method is not supported

Source code in holocron/nn/modules/loss.py

def __init__(
    self,
    weight: float | list[float] | Tensor | None = None,
    ignore_index: int = -100,
    reduction: str = "mean",
) -> None:
    super().__init__()
    # Cast class weights if possible
    self.weight: Tensor | None
    if isinstance(weight, (float, int)):
        self.register_buffer("weight", torch.Tensor([weight, 1 - weight]))
    elif isinstance(weight, list):
        self.register_buffer("weight", torch.Tensor(weight))
    elif isinstance(weight, Tensor):
        self.register_buffer("weight", weight)
    else:
        self.weight: Tensor | None = None
    self.ignore_index: int = ignore_index
    # Set the reduction method
    if reduction not in {"none", "mean", "sum"}:
        raise NotImplementedError("argument reduction received an incorrect input")
    self.reduction: str = reduction

FocalLoss

FocalLoss(gamma: float = 2.0, **kwargs: Any)

Bases: Loss

Implementation of Focal Loss as described in "Focal Loss for Dense Object Detection".

While the weighted cross-entropy is described by:

\[ CE(p_t) = -\alpha_t log(p_t) \]

where \(\alpha_t\) is the loss weight of class \(t\), and \(p_t\) is the predicted probability of class \(t\).

the focal loss introduces a modulating factor

\[ FL(p_t) = -\alpha_t (1 - p_t)^\gamma log(p_t) \]

where \(\gamma\) is a positive focusing parameter.

PARAMETER	DESCRIPTION
gamma	exponent parameter of the focal loss TYPE: float DEFAULT: 2.0
**kwargs	keyword args of Loss TYPE: Any DEFAULT: {}

Source code in holocron/nn/modules/loss.py

def __init__(self, gamma: float = 2.0, **kwargs: Any) -> None:
    super().__init__(**kwargs)
    self.gamma: float = gamma

MultiLabelCrossEntropy

MultiLabelCrossEntropy(*args: Any, **kwargs: Any)

Bases: Loss

Implementation of the cross-entropy loss for multi-label targets

PARAMETER	DESCRIPTION
*args	args of Loss TYPE: Any DEFAULT: ()
**kwargs	keyword args of Loss TYPE: Any DEFAULT: {}

Source code in holocron/nn/modules/loss.py

def __init__(self, *args: Any, **kwargs: Any) -> None:
    super().__init__(*args, **kwargs)

ComplementCrossEntropy

ComplementCrossEntropy(gamma: float = -1, **kwargs: Any)

Bases: Loss

Implements the complement cross entropy loss from "Imbalanced Image Classification with Complement Cross Entropy"

PARAMETER	DESCRIPTION
gamma	smoothing factor TYPE: float DEFAULT: -1
**kwargs	keyword args of Loss TYPE: Any DEFAULT: {}

Source code in holocron/nn/modules/loss.py

def __init__(self, gamma: float = -1, **kwargs: Any) -> None:
    super().__init__(**kwargs)
    self.gamma: float = gamma

MutualChannelLoss

MutualChannelLoss(weight: float | list[float] | Tensor | None = None, ignore_index: int = -100, reduction: str = 'mean', xi: int = 2, alpha: float = 1)

Bases: Loss

Implements the mutual channel loss from "The Devil is in the Channels: Mutual-Channel Loss for Fine-Grained Image Classification".

PARAMETER	DESCRIPTION
weight	class weight for loss computation TYPE: float | list[float] | Tensor | None DEFAULT: None
ignore_index	specifies target value that is ignored and do not contribute to gradient TYPE: int DEFAULT: -100
reduction	type of reduction to apply to the final loss TYPE: str DEFAULT: 'mean'
xi	num of features per class TYPE: int DEFAULT: 2
alpha	diversity factor TYPE: float DEFAULT: 1

Source code in holocron/nn/modules/loss.py

def __init__(
    self,
    weight: float | list[float] | Tensor | None = None,
    ignore_index: int = -100,
    reduction: str = "mean",
    xi: int = 2,
    alpha: float = 1,
) -> None:
    super().__init__(weight, ignore_index, reduction)
    self.xi: int = xi
    self.alpha: float = alpha

DiceLoss

DiceLoss(weight: float | list[float] | Tensor | None = None, gamma: float = 1.0, eps: float = 1e-08)

Bases: Loss

Implements the dice loss from "V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation".

PARAMETER	DESCRIPTION
weight	class weight for loss computation TYPE: float | list[float] | Tensor | None DEFAULT: None
gamma	recall/precision control param TYPE: float DEFAULT: 1.0
eps	small value added to avoid division by zero TYPE: float DEFAULT: 1e-08

Source code in holocron/nn/modules/loss.py

def __init__(
    self,
    weight: float | list[float] | Tensor | None = None,
    gamma: float = 1.0,
    eps: float = 1e-8,
) -> None:
    super().__init__(weight)
    self.gamma: float = gamma
    self.eps: float = eps

PolyLoss

PolyLoss(*args: Any, eps: float = 2.0, **kwargs: Any)

Bases: Loss

Implements the Poly1 loss from "PolyLoss: A Polynomial Expansion Perspective of Classification Loss Functions".

Pass raw logits with shape (N, K, ...), where K is the number of classes. Hard targets have shape (N, ...) and dtype torch.int64, with values in [0, K). Soft targets are floating-point class probabilities with the same shape as the logits. Do not apply softmax to the input first.

Example

import torch from holocron.nn import PolyLoss logits = torch.randn(4, 10, requires_grad=True) target = torch.tensor([0, 9, 3, 1], dtype=torch.int64) loss = PolyLoss()(logits, target) loss.backward()

Note

Out-of-range hard targets, including the default ignore_index=-100, raise RuntimeError if present in the target. Only an in-range class index can be ignored for sum/mean reduction. With reduction="none", ignored positions are not zeroed and hard-target losses are returned as a flat tensor. See the loss input guide for a loss that supports out-of-range ignored targets.

PARAMETER	DESCRIPTION
*args	positional weight, ignore_index and reduction arguments of Loss TYPE: Any DEFAULT: ()
eps	epsilon 1 from the paper (default: 2.0) TYPE: float DEFAULT: 2.0
**kwargs	weight, ignore_index and reduction keyword arguments of Loss TYPE: Any DEFAULT: {}

Source code in holocron/nn/modules/loss.py

def __init__(
    self,
    *args: Any,
    eps: float = 2.0,
    **kwargs: Any,
) -> None:
    super().__init__(*args, **kwargs)
    self.eps: float = eps

Loss wrappers

ClassBalancedWrapper

ClassBalancedWrapper(criterion: Module, num_samples: Tensor, beta: float = 0.99)

Bases: Module

Implementation of the class-balanced loss as described in "Class-Balanced Loss Based on Effective Number of Samples".

Given a loss function \(\mathcal{L}\), the class-balanced loss is described by:

\[ CB(p, y) = \frac{1 - \beta}{1 - \beta^{n_y}} \mathcal{L}(p, y) \]

where \(p\) is the predicted probability for class \(y\), \(n_y\) is the number of training samples for class \(y\), and \(\beta\) is exponential factor.

PARAMETER	DESCRIPTION
criterion	loss module TYPE: Module
num_samples	number of samples for each class TYPE: Tensor
beta	rebalancing exponent TYPE: float DEFAULT: 0.99

Source code in holocron/nn/modules/loss.py

def __init__(self, criterion: nn.Module, num_samples: Tensor, beta: float = 0.99) -> None:
    super().__init__()
    self.criterion = criterion
    self.beta: float = beta
    cb_weights = (1 - beta) / (1 - beta**num_samples)
    if self.criterion.weight is None:
        self.criterion.weight: Tensor | None = cb_weights
    else:
        self.criterion.weight *= cb_weights.to(device=self.criterion.weight.device)  # ty: ignore[invalid-argument-type]

Convolution layers

NormConv2d

NormConv2d(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, padding_mode: Literal['zeros', 'reflect', 'replicate', 'circular'] = 'zeros', eps: float = 1e-14)

Bases: _NormConvNd

Implements the normalized convolution module from "Normalized Convolutional Neural Network".

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

\[ out(N_i, C_{out_j}) = bias(C_{out_j}) + \sum_{k = 0}^{C_{in} - 1} weight(C_{out_j}, k) \star \frac{input(N_i, k) - \mu(N_i, k)}{\sqrt{\sigma^2(N_i, k) + \epsilon}} \]

where: \(\star\) is the valid 2D cross-correlation operator, \(\mu(N_i, k)\) and \(\sigma²(N_i, k)\) are the mean and variance of \(input(N_i, k)\) over all slices, \(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.

PARAMETER	DESCRIPTION
in_channels	Number of channels in the input image TYPE: int
out_channels	Number of channels produced by the convolution TYPE: int
kernel_size	Size of the convolving kernel TYPE: int
stride	Stride of the convolution. TYPE: int DEFAULT: 1
padding	Zero-padding added to both sides of the input. TYPE: int DEFAULT: 0
dilation	Spacing between kernel elements. TYPE: int DEFAULT: 1
groups	Number of blocked connections from input channels to output channels. TYPE: int DEFAULT: 1
bias	If True, adds a learnable bias to the output. TYPE: bool DEFAULT: True
padding_mode	padding mode for the convolution. TYPE: Literal['zeros', 'reflect', 'replicate', 'circular'] DEFAULT: 'zeros'
eps	a value added to the denominator for numerical stability. TYPE: float DEFAULT: 1e-14

Source code in holocron/nn/modules/conv.py

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
    eps: float = 1e-14,
) -> None:
    kernel_size = _pair(kernel_size)
    stride = _pair(stride)
    padding = _pair(padding)
    dilation = _pair(dilation)
    super().__init__(
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        dilation,
        False,
        _pair(0),
        groups,
        bias,
        padding_mode,
        False,
        eps,
    )

Add2d

Add2d(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, padding_mode: Literal['zeros', 'reflect', 'replicate', 'circular'] = 'zeros', normalize_slices: bool = False, eps: float = 1e-14)

Bases: _NormConvNd

Implements the adder module from "AdderNet: Do We Really Need Multiplications in Deep Learning?".

In the simplest case, the output value of the layer at position \((m, n)\) in channel \(c\) with filter F of spatial size \((d, d)\), intput size \((C_{in}, H, W)\) and output \((C_{out}, H, W)\) can be precisely described as:

\[ out(m, n, c) = - \sum\limits_{i=0}^d \sum\limits_{j=0}^d \sum\limits_{k=0}^{C_{in}} |X(m + i, n + j, k) - F(i, j, k, c)| \]

where \(C\) denotes a number of channels, \(H\) is a height of input planes in pixels, and \(W\) is width in pixels.

PARAMETER	DESCRIPTION
in_channels	Number of channels in the input image TYPE: int
out_channels	Number of channels produced by the convolution TYPE: int
kernel_size	Size of the convolving kernel TYPE: int
stride	Stride of the convolution. TYPE: int DEFAULT: 1
padding	Zero-padding added to both sides of the input. TYPE: int DEFAULT: 0
dilation	Spacing between kernel elements. TYPE: int DEFAULT: 1
groups	Number of blocked connections from input channels to output channels. TYPE: int DEFAULT: 1
bias	If True, adds a learnable bias to the output. TYPE: bool DEFAULT: True
padding_mode	padding mode for the convolution. TYPE: Literal['zeros', 'reflect', 'replicate', 'circular'] DEFAULT: 'zeros'
normalize_slices	whether slices should be normalized before performing cross-correlation. TYPE: bool DEFAULT: False
eps	a value added to the denominator for numerical stability. TYPE: float DEFAULT: 1e-14

Source code in holocron/nn/modules/conv.py

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
    normalize_slices: bool = False,
    eps: float = 1e-14,
) -> None:
    kernel_size = _pair(kernel_size)
    stride = _pair(stride)
    padding = _pair(padding)
    dilation = _pair(dilation)
    super().__init__(
        in_channels,
        out_channels,
        kernel_size,
        stride,
        padding,
        dilation,
        False,
        _pair(0),
        groups,
        bias,
        padding_mode,
        normalize_slices,
        eps,
    )

SlimConv2d

SlimConv2d(in_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, padding_mode: Literal['zeros', 'reflect', 'replicate', 'circular'] = 'zeros', r: int = 32, L: int = 2)

Bases: Module

Implements the convolution module from "SlimConv: Reducing Channel Redundancy in Convolutional Neural Networks by Weights Flipping".

First, we compute channel-wise weights as follows:

\[ z(c) = \frac{1}{H \cdot W} \sum\limits_{i=1}^H \sum\limits_{j=1}^W X_{c,i,j} \]

where \(X \in \mathbb{R}^{C \times H \times W}\) is the input tensor, \(H\) is height in pixels, and \(W\) is width in pixels.

\[ w = \sigma(F_{fc2}(\delta(F_{fc1}(z)))) \]

where \(z \in \mathbb{R}^{C}\) contains channel-wise statistics, \(\sigma\) refers to the sigmoid function, \(\delta\) refers to the ReLU function, \(F_{fc1}\) is a convolution operation with kernel of size \((1, 1)\) with \(max(C/r, L)\) output channels followed by batch normalization, and \(F_{fc2}\) is a plain convolution operation with kernel of size \((1, 1)\) with \(C\) output channels.

We then proceed with reconstructing and transforming both pathways:

\[ X_{top} = X \odot w X_{bot} = X \odot \check{w} \]

where \(\odot\) refers to the element-wise multiplication and \(\check{w}\) is the channel-wise reverse-flip of \(w\).

\[ T_{top} = F_{top}(X_{top}^{(1)} + X_{top}^{(2)}) T_{bot} = F_{bot}(X_{bot}^{(1)} + X_{bot}^{(2)}) \]

where \(X^{(1)}\) and \(X^{(2)}\) are the channel-wise first and second halves of \(X\), \(F_{top}\) is a convolution of kernel size \((3, 3)\), and \(F_{bot}\) is a convolution of kernel size \((1, 1)\) reducing channels by half, followed by a convolution of kernel size \((3, 3)\).

Finally we fuse both pathways to yield the output:

\[ Y = T_{top} \oplus T_{bot} \]

where \(\oplus\) is the channel-wise concatenation.

PARAMETER	DESCRIPTION
in_channels	Number of channels in the input image TYPE: int
kernel_size	Size of the convolving kernel TYPE: int
stride	Stride of the convolution. TYPE: int DEFAULT: 1
padding	Zero-padding added to both sides of the input. TYPE: int DEFAULT: 0
dilation	Spacing between kernel elements. TYPE: int DEFAULT: 1
groups	Number of blocked connections from input channels to output channels. TYPE: int DEFAULT: 1
bias	If True, adds a learnable bias to the output. TYPE: bool DEFAULT: True
padding_mode	padding mode for the convolution. TYPE: Literal['zeros', 'reflect', 'replicate', 'circular'] DEFAULT: 'zeros'
r	squeezing divider. TYPE: int DEFAULT: 32
L	minimum squeezed channels. TYPE: int DEFAULT: 2

Source code in holocron/nn/modules/conv.py

def __init__(
    self,
    in_channels: int,
    kernel_size: int,
    stride: int = 1,
    padding: int = 0,
    dilation: int = 1,
    groups: int = 1,
    bias: bool = True,
    padding_mode: Literal["zeros", "reflect", "replicate", "circular"] = "zeros",
    r: int = 32,
    L: int = 2,  # noqa: N803
) -> None:
    super().__init__()
    self.fc1 = nn.Conv2d(in_channels, max(in_channels // r, L), 1)
    self.bn = nn.BatchNorm2d(max(in_channels // r, L))
    self.fc2 = nn.Conv2d(max(in_channels // r, L), in_channels, 1)
    self.conv_top = nn.Conv2d(
        in_channels // 2, in_channels // 2, kernel_size, stride, padding, dilation, groups, bias, padding_mode
    )
    self.conv_bot1 = nn.Conv2d(in_channels // 2, in_channels // 4, 1)
    self.conv_bot2 = nn.Conv2d(
        in_channels // 4, in_channels // 4, kernel_size, stride, padding, dilation, groups, bias, padding_mode
    )

PyConv2d

PyConv2d(in_channels: int, out_channels: int, kernel_size: int, num_levels: int = 2, padding: int = 0, groups: list[int] | None = None, **kwargs: Any)

Bases: ModuleList

Implements the convolution module from "Pyramidal Convolution: Rethinking Convolutional Neural Networks for Visual Recognition".

PARAMETER	DESCRIPTION
in_channels	Number of channels in the input image TYPE: int
out_channels	Number of channels produced by the convolution TYPE: int
kernel_size	Size of the convolving kernel TYPE: int
num_levels	number of stacks in the pyramid. TYPE: int DEFAULT: 2
padding	Zero-padding added to both sides of the input. TYPE: int DEFAULT: 0
groups	Number of blocked connections from input channels to output channels. TYPE: list[int] | None DEFAULT: None
kwargs	keyword args of torch.nn.Conv2d. TYPE: Any DEFAULT: {}

Source code in holocron/nn/modules/conv.py

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int,
    num_levels: int = 2,
    padding: int = 0,
    groups: list[int] | None = None,
    **kwargs: Any,
) -> None:
    if num_levels == 1:
        super().__init__([
            nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size,
                padding=padding,
                groups=groups[0] if isinstance(groups, list) else 1,
                **kwargs,
            )
        ])
    else:
        exp2 = int(math.log2(num_levels))
        reminder = num_levels - 2**exp2
        out_chans = [out_channels // 2 ** (exp2 + 1)] * (2 * reminder) + [out_channels // 2**exp2] * (
            num_levels - 2 * reminder
        )

        k_sizes = [kernel_size + 2 * idx for idx in range(num_levels)]
        if groups is None:
            groups = [1] + [
                min(2 ** (2 + idx), out_chan)
                for idx, out_chan in zip(range(num_levels - 1), out_chans[1:], strict=True)
            ]
        elif not isinstance(groups, list) or len(groups) != num_levels:
            raise ValueError("The argument `group` is expected to be a list of integer of size `num_levels`.")
        paddings = [padding + idx for idx in range(num_levels)]

        super().__init__([
            nn.Conv2d(in_channels, out_chan, k_size, padding=padding, groups=group, **kwargs)
            for out_chan, k_size, padding, group in zip(out_chans, k_sizes, paddings, groups, strict=True)
        ])
    self.num_levels: int = num_levels

Involution2d

Involution2d(in_channels: int, kernel_size: int, padding: int = 0, stride: int = 1, groups: int = 1, dilation: int = 1, reduction_ratio: float = 1)

Bases: Module

Implements the convolution module from "Involution: Inverting the Inherence of Convolution for Visual Recognition", adapted from the proposed PyTorch implementation in the paper.

PARAMETER	DESCRIPTION
in_channels	Number of channels in the input image TYPE: int
kernel_size	Size of the convolving kernel TYPE: int
padding	Zero-padding added to both sides of the input. TYPE: int DEFAULT: 0
stride	Stride of the convolution. TYPE: int DEFAULT: 1
groups	Number of blocked connections from input channels to output channels. TYPE: int DEFAULT: 1
dilation	Spacing between kernel elements. TYPE: int DEFAULT: 1
reduction_ratio	reduction ratio of the channels to generate the kernel TYPE: float DEFAULT: 1

Source code in holocron/nn/modules/conv.py

def __init__(
    self,
    in_channels: int,
    kernel_size: int,
    padding: int = 0,
    stride: int = 1,
    groups: int = 1,
    dilation: int = 1,
    reduction_ratio: float = 1,
) -> None:
    super().__init__()

    self.groups: int = groups
    self.k_size: int = kernel_size

    self.pool: nn.AvgPool2d | None = nn.AvgPool2d(stride, stride) if stride > 1 else None
    self.reduce = nn.Conv2d(in_channels, int(in_channels // reduction_ratio), 1)
    self.span = nn.Conv2d(int(in_channels // reduction_ratio), kernel_size**2 * groups, 1)
    self.unfold = nn.Unfold(kernel_size, dilation, padding, stride)

Regularization layers

DropBlock2d

DropBlock2d(p: float = 0.1, block_size: int = 7, inplace: bool = False)

Bases: Module

Implements the DropBlock module from "DropBlock: A regularization method for convolutional networks"

PARAMETER	DESCRIPTION
p	probability of dropping activation value TYPE: float DEFAULT: 0.1
block_size	size of each block that is expended from the sampled mask TYPE: int DEFAULT: 7
inplace	whether the operation should be done inplace TYPE: bool DEFAULT: False

Source code in holocron/nn/modules/dropblock.py

def __init__(self, p: float = 0.1, block_size: int = 7, inplace: bool = False) -> None:
    super().__init__()
    self.p: float = p
    self.block_size: int = block_size
    self.inplace: bool = inplace

Downsampling

ConcatDownsample2d

ConcatDownsample2d(scale_factor: int)

Bases: Module

Implements a loss-less downsampling operation described in "YOLO9000: Better, Faster, Stronger" by stacking adjacent information on the channel dimension.

PARAMETER	DESCRIPTION
scale_factor	spatial scaling factor TYPE: int

Source code in holocron/nn/modules/downsample.py

def __init__(self, scale_factor: int) -> None:
    super().__init__()
    self.scale_factor: int = scale_factor

GlobalAvgPool2d

GlobalAvgPool2d(flatten: bool = False)

Bases: Module

Fast implementation of global average pooling from "TResNet: High Performance GPU-Dedicated Architecture"

PARAMETER	DESCRIPTION
flatten	whether spatial dimensions should be squeezed TYPE: bool DEFAULT: False

Source code in holocron/nn/modules/downsample.py

def __init__(self, flatten: bool = False) -> None:
    super().__init__()
    self.flatten: bool = flatten

GlobalMaxPool2d

GlobalMaxPool2d(flatten: bool = False)

Bases: Module

Fast implementation of global max pooling from "TResNet: High Performance GPU-Dedicated Architecture"

PARAMETER	DESCRIPTION
flatten	whether spatial dimensions should be squeezed TYPE: bool DEFAULT: False

Source code in holocron/nn/modules/downsample.py

def __init__(self, flatten: bool = False) -> None:
    super().__init__()
    self.flatten: bool = flatten

BlurPool2d

BlurPool2d(channels: int, kernel_size: int = 3, stride: int = 2)

Bases: Module

Ross Wightman's implementation of blur pooling module as described in "Making Convolutional Networks Shift-Invariant Again".

PARAMETER	DESCRIPTION
channels	Number of input channels TYPE: int
kernel_size	binomial filter size for blurring. currently supports 3 (default) and 5. TYPE: int DEFAULT: 3
stride	downsampling filter stride TYPE: int DEFAULT: 2

Source code in holocron/nn/modules/downsample.py

def __init__(self, channels: int, kernel_size: int = 3, stride: int = 2) -> None:
    super().__init__()
    self.channels: int = channels
    if kernel_size <= 1:
        raise AssertionError
    self.kernel_size: int = kernel_size
    self.stride: int = stride
    pad_size = [get_padding(kernel_size, stride, dilation=1)] * 4
    self.padding = nn.ReflectionPad2d(pad_size)  # type: ignore[arg-type]
    self._coeffs = torch.tensor((np.poly1d((0.5, 0.5)) ** (self.kernel_size - 1)).coeffs)  # for torchscript compat
    self.kernel: dict[str, Tensor] = {}  # lazy init by device for DataParallel compat

SPP

SPP(kernel_sizes: list[int])

Bases: ModuleList

SPP layer from "Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition".

PARAMETER	DESCRIPTION
kernel_sizes	kernel sizes of each pooling TYPE: list[int]

Source code in holocron/nn/modules/downsample.py

def __init__(self, kernel_sizes: list[int]) -> None:
    super().__init__([nn.MaxPool2d(k_size, stride=1, padding=k_size // 2) for k_size in kernel_sizes])

ZPool

ZPool(dim: int = 1)

Bases: Module

Z-pool layer from "Rotate to Attend: Convolutional Triplet Attention Module".

PARAMETER	DESCRIPTION
dim	dimension to pool across TYPE: int DEFAULT: 1

Source code in holocron/nn/modules/downsample.py

def __init__(self, dim: int = 1) -> None:
    super().__init__()
    self.dim: int = dim

Attention

SAM

SAM(in_channels: int)

Bases: Module

SAM layer from "CBAM: Convolutional Block Attention Module" modified in "YOLOv4: Optimal Speed and Accuracy of Object Detection".

PARAMETER	DESCRIPTION
in_channels	input channels TYPE: int

Source code in holocron/nn/modules/attention.py

def __init__(self, in_channels: int) -> None:
    super().__init__()
    self.conv = nn.Conv2d(in_channels, 1, 1)

LambdaLayer

LambdaLayer(in_channels: int, out_channels: int, dim_k: int, n: int | None = None, r: int | None = None, num_heads: int = 4, dim_u: int = 1)

Bases: Module

Lambda layer from "LambdaNetworks: Modeling long-range interactions without attention". The implementation was adapted from lucidrains.

PARAMETER	DESCRIPTION
in_channels	input channels TYPE: int
out_channels	output channels TYPE: int
dim_k	key dimension TYPE: int
n	number of input pixels TYPE: int | None DEFAULT: None
r	receptive field for relative positional encoding TYPE: int | None DEFAULT: None
num_heads	number of attention heads TYPE: int DEFAULT: 4
dim_u	intra-depth dimension TYPE: int DEFAULT: 1

Source code in holocron/nn/modules/lambda_layer.py

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    dim_k: int,
    n: int | None = None,
    r: int | None = None,
    num_heads: int = 4,
    dim_u: int = 1,
) -> None:
    super().__init__()
    self.u: int = dim_u
    self.num_heads: int = num_heads

    if out_channels % num_heads != 0:
        raise AssertionError("values dimension must be divisible by number of heads for multi-head query")
    dim_v = out_channels // num_heads

    # Project input and context to get queries, keys & values
    self.to_q = nn.Conv2d(in_channels, dim_k * num_heads, 1, bias=False)
    self.to_k = nn.Conv2d(in_channels, dim_k * dim_u, 1, bias=False)
    self.to_v = nn.Conv2d(in_channels, dim_v * dim_u, 1, bias=False)

    self.norm_q = nn.BatchNorm2d(dim_k * num_heads)
    self.norm_v = nn.BatchNorm2d(dim_v * dim_u)

    self.local_contexts: bool = r is not None
    if r is not None:
        if r % 2 != 1:
            raise AssertionError("Receptive kernel size should be odd")
        self.padding: int = r // 2
        self.R = nn.Parameter(torch.randn(dim_k, dim_u, 1, r, r))
    else:
        if n is None:
            raise AssertionError("You must specify the total sequence length (h x w)")
        self.pos_emb = nn.Parameter(torch.randn(n, n, dim_k, dim_u))

TripletAttention

TripletAttention()

Bases: Module

Triplet attention layer from "Rotate to Attend: Convolutional Triplet Attention Module". This implementation is based on the one from the paper's authors.

Source code in holocron/nn/modules/attention.py

def __init__(self) -> None:
    super().__init__()
    self.c_branch = DimAttention(dim=1)
    self.h_branch = DimAttention(dim=2)
    self.w_branch = DimAttention(dim=3)