Convolutional 2D (Conv2D)#

API#

class warp_nn.modules.layers.Conv2D(
in_channels: int,
out_channels: int,
kernel_size: int | tuple[int, int],
*,
stride: int | tuple[int, int] = 1,
padding: int | tuple[int, int] = 0,
dilation: int | tuple[int, int] = 1,
groups: int = 1,
bias: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply a 2D convolution.

\[\text{Conv2D}(x)_{i, j} = b_j + \sum_{c = 0}^{C_{in} - 1} W_{j, c} \star x_{i, c}\]

where \(\star\) is the valid 2D cross-correlation operator, \(i\) indexes the batch dimension, \(j\) the output channel, \(c\) the input channel (for groups = 1), and \(C_{in}\) is in_channels.

Learnable parameters:

Name

Shape

Description

\(W\)

weight

(out_channels, in_channels / groups, kernel_size[0], kernel_size[1])

Weights

\(b\)

bias

(out_channels, 1)

Bias. Only if bias is true

The parameters are initialized from the uniform distribution \(u(-k, k)\) where \(k = \sqrt{\frac{groups}{\text{in\_channels * kernel\_size[0] * kernel\_size[1]}}}\)

Parameters:

  • in_channels – The number of input channels.

  • out_channels – The number of output channels.

  • kernel_size – The size of the kernel.

  • stride – The stride of the convolution.

  • padding – The padding of the convolution.

  • dilation – The dilation of the convolution.

  • groups – The number of groups. Both, the in_channels and the out_channels arguments must be divisible by groups.

  • bias – Whether to include a bias term.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

__call__(
input: array,
) → array[source]#

Forward pass of the module.

Parameters:

input – The input array, with shape (batch_size, in_channels, in_height, in_width).

Returns:

The output array, with shape (batch_size, out_channels, out_height, out_width) where:

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