Convolution Operators

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Forward convolution

tileops.convolution.Conv1dFwdOp

1D convolution over an NCL input, as torch.nn.functional.conv1d.

__init__

__init__(
    stride=1,
    padding=0,
    dilation=1,
    groups=1,
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from the first call.

Parameters:

• stride (int | Tuple[int], default: 1 ) –

  Stride, an int or a 1-tuple (default 1).

• padding (int | Tuple[int] | str, default: 0 ) –

  Padding each side, an int or a 1-tuple, or 'valid' / 'same' (default 0). 'same' puts the extra element of an odd total on the right.

• dilation (int | Tuple[int], default: 1 ) –

  Dilation, an int or a 1-tuple (default 1).

• groups (int, default: 1 ) –

  Number of channel groups (default 1).

• target (Target, default: None ) –

  Backend target to serve this op, or None to decide from the input device.

forward

forward(
    input,
    weight,
    bias=None,
)

Apply the convolution.

Parameters:

• input (Tensor) –

  Input, \([N \times C_{in} \times L]\).

• weight (Tensor) –

  Weight, \([C_{out} \times C_{in}/\text{groups} \times k]\).

• bias (Optional[Tensor], default: None ) –

  Per-output-channel bias, \([C_{out}]\), or None.

Returns:

• Tensor –

  The convolution result, \([N \times C_{out} \times L_{out}]\).

tileops.convolution.Conv2dFwdOp

2D convolution over an NCHW input, as torch.nn.functional.conv2d.

They multiply float32 operands in TF32, as cuDNN does by default.

__init__

__init__(
    stride=1,
    padding=0,
    dilation=1,
    groups=1,
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from the first call.

Parameters:

• stride (int | Tuple[int, int], default: 1 ) –

  Stride, an int or a 2-tuple (default 1).

• padding (int | Tuple[int, int] | str, default: 0 ) –

  Padding each side, an int or a 2-tuple, or 'valid' / 'same' (default 0). 'same' puts the extra element of an odd total after.

• dilation (int | Tuple[int, int], default: 1 ) –

  Dilation, an int or a 2-tuple (default 1).

• groups (int, default: 1 ) –

  Number of channel groups (default 1).

• target (Target, default: None ) –

  Backend target to serve this op, or None to decide from the input device.

forward

forward(
    input,
    weight,
    bias=None,
)

Apply the convolution.

Parameters:

• input (Tensor) –

  Input, \([N \times C_{in} \times H \times W]\).

• weight (Tensor) –

  Weight, \([C_{out} \times C_{in}/\text{groups} \times k_H \times k_W]\).

• bias (Optional[Tensor], default: None ) –

  Per-output-channel bias, \([C_{out}]\), or None.

Returns:

• Tensor –

  The convolution result, \([N \times C_{out} \times H_{out} \times W_{out}]\).

tileops.convolution.Conv3dFwdOp

3D convolution over an NCDHW input, as torch.nn.functional.conv3d.

They multiply float32 operands in TF32, as cuDNN does by default.

__init__

__init__(
    stride=1,
    padding=0,
    dilation=1,
    groups=1,
    *,
    target=None
)

Build the op. Shapes and dtypes are taken from the first call.

Parameters:

• stride (int | Tuple[int, int, int], default: 1 ) –

  Stride, an int or a 3-tuple (default 1).

• padding (int | Tuple[int, int, int] | str, default: 0 ) –

  Padding each side, an int or a 3-tuple, or 'valid' / 'same' (default 0). 'same' puts the extra element of an odd total after.

• dilation (int | Tuple[int, int, int], default: 1 ) –

  Dilation, an int or a 3-tuple (default 1).

• groups (int, default: 1 ) –

  Number of channel groups (default 1).

• target (Target, default: None ) –

  Backend target to serve this op, or None to decide from the input device.

forward

forward(
    input,
    weight,
    bias=None,
)

Apply the convolution.

Parameters:

• input (Tensor) –

  Input, \([N \times C_{in} \times D \times H \times W]\).

• weight (Tensor) –

  Weight, \([C_{out} \times C_{in}/\text{groups} \times k_D \times k_H \times k_W]\).

• bias (Optional[Tensor], default: None ) –

  Per-output-channel bias, \([C_{out}]\), or None.

Returns:

• Tensor –

  The convolution result, $[N \times C_{out} \times D_{out} \times H_{out}

• Tensor –

  \times W_{out}]$.