Skip to main content

/

Docs

Nightly
v24.5
Log in

Mojo function

conv2d

conv2d(input: Symbol, filter: Symbol, stride: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, 1)), dilation: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, 1)), padding: Tuple[Int, Int, Int, Int] = Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0}), store_to_mem({0}), store_to_mem({0})>, 1)), groups: Int = 1) -> Symbol

Computes the 2-D convolution product of the input with the given filter, strides, dilations, paddings, and groups.

The op supports 2-D convolution, with the following layout assumptions:

  • input has NHWC layout, i.e., (batch_size, height, width, in_channels)
  • filter has layout RSCF, i.e., (height, width, in_channels / num_groups, out_channels)

The padding values are expected to take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after...) and represent padding 0's before and after the indicated spatial dimensions in input. In 2-D convolution, dim1 here repesents H and dim2 represents W. In Python like syntax, padding a 2x3 spatial input with [0, 1, 2, 1] would yield:

input = [
  [1, 2, 3],
  [4, 5, 6]
]
# Shape is 2x3

padded_input = [
  [0, 0, 1, 2, 3, 0],
  [0, 0, 4, 5, 6, 0]
  [0, 0, 0, 0, 0, 0]
]
# Shape is 3x6
input = [
  [1, 2, 3],
  [4, 5, 6]
]
# Shape is 2x3

padded_input = [
  [0, 0, 1, 2, 3, 0],
  [0, 0, 4, 5, 6, 0]
  [0, 0, 0, 0, 0, 0]
]
# Shape is 3x6

This op currently only supports strides and padding on the input.

Args:

  • ​input (Symbol): An NHWC input tensor to perform the convolution upon.
  • ​filter (Symbol): The convolution filter in RSCF layout: (height, width, in_channels / num_groups, out_channels).
  • ​stride (Tuple[Int, Int]): The stride of the convolution operation.
  • ​dilation (Tuple[Int, Int]): The spacing between the kernel points.
  • ​padding (Tuple[Int, Int, Int, Int]): The amount of padding applied to the input.
  • ​groups (Int): When greater than 1, divides the convolution into multiple parallel convolutions. The number of input and output channels must both be divisible by the number of groups.

Returns:

A symbolic tensor value with the convolution applied.