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,
)[source]#

Bases: Module

Apply a 2D convolution.

\[\text{Conv2D}(x) = TODO\]

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.

__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} \]