Skip to main content

Python class

Conv2d

Conv2d

class max.nn.module_v3.Conv2d(kernel_size, in_channels, out_channels, dtype=None, stride=1, padding=0, dilation=1, num_groups=1, device=None, has_bias=False, permute=False, name=None)

A 2D convolution layer.

This is a Conv2d implementation that uses Tensor instead of Weight objects.

Example:

from max.nn.module_v3 import Conv2d
from max.experimental.tensor import Tensor

conv = Conv2d(
    kernel_size=3,
    in_channels=3,
    out_channels=64,
    has_bias=True,
    permute=True,
)

x = Tensor.ones([1, 3, 32, 32])
result = conv(x)

Initialize Conv2d layer.

Parameters:

  • kernel_size (int | tuple[int, int]) – Size of the convolving kernel. Can be a single int (square kernel) or tuple (height, width).
  • in_channels (int) – Number of channels in the input image.
  • out_channels (int) – Number of channels produced by the convolution.
  • dtype (DType | None) – The data type for both weights and bias. In v3, this is optional as Tensor manages dtype automatically.
  • stride (int | tuple[int, int]) – Stride of the convolution for height and width dimensions. Can be int (applied to both dimensions) or tuple (stride_h, stride_w). Default: 1
  • padding (int | tuple[int, int] | tuple[int, int, int, int]) – Padding added to input. Can be int (applied to all sides), tuple of 2 ints (pad_h, pad_w), or tuple of 4 ints (pad_top, pad_bottom, pad_left, pad_right) to support asymmetric padding. Default: 0
  • dilation (int | tuple[int, int]) – Spacing between kernel elements for height and width dimensions. Can be int (applied to both dimensions) or tuple (dilation_h, dilation_w). Default: 1
  • num_groups (int) – Number of blocked connections from input channels to output channels. Input channels and output channels are divided into groups. Default: 1
  • device (DeviceRef | None) – The target device for computation. In v3, this is optional as Tensor manages device automatically.
  • has_bias (bool) – If true, adds a learnable bias vector to the layer. Defaults to False.
  • permute (bool) – If true, permutes weights from PyTorch format to MAX format. PyTorch order: (out_channels, in_channels / num_groups, height, width). MAX API order: (height, width, in_channels / num_groups, out_channels). Defaults to False.
  • name (str | None) – Base name for weights. In v3, this is stored but not used for Weight naming. Defaults to None.

bias

bias: Tensor | Literal[0]

The bias tensor with shape [out_channels] (or 0 if bias is disabled).

forward()

forward(x)

Apply 2D convolution to input.

Parameters:

x (Tensor) – Input tensor. Shape depends on permute:

  • If permute=True: [batch_size, in_channels, height, width]
  • If permute=False: [batch_size, height, width, in_channels]

Returns:

  • If permute=True: [batch_size, out_channels, new_height, new_width]
  • If permute=False: [batch_size, new_height, new_width, out_channels]

Return type:

Output tensor. Shape depends on permute

weight

weight: Tensor

The weight tensor with shape [out_channels, in_channels // num_groups, kernel_height, kernel_width].

Was this page helpful?