mlx.nn.Conv2d

Contents

mlx.nn.Conv2d#

class Conv2d(in_channels: int, out_channels: int, kernel_size: int | tuple, stride: int | tuple = 1, padding: int | tuple = 0, dilation: int | tuple = 1, groups: int = 1, bias: bool = True)#

Applies a 2-dimensional convolution over the multi-channel input image.

The channels are expected to be last i.e. the input shape should be NHWC where:

  • N is the batch dimension

  • H is the input image height

  • W is the input image width

  • C is the number of input channels

Parameters:

  • in_channels (int) – The number of input channels.

  • out_channels (int) – The number of output channels.

  • kernel_size (int or tuple) – The size of the convolution filters.

  • stride (int or tuple, optional) – The size of the stride when applying the filter. Default: 1.

  • padding (int or tuple, optional) – How many positions to 0-pad the input with. Default: 0.

  • dilation (int or tuple, optional) – The dilation of the convolution.

  • groups (int, optional) – The number of groups for the convolution. Default: 1.

  • bias (bool, optional) – If True add a learnable bias to the output. Default: True

Methods

Contents