avgpool¶

tripy.avgpool(input: Tensor, kernel_dims: Sequence[int], stride: Sequence[int] | None = None, padding: Sequence[Tuple[int]] | None = None) → Tensor[source]¶

Applies an average pooling over the input tensor.

The output’s non-spatial dimensions are the same as input. For each input spatial dimension \(D_{i}\), the corresponding output dimension will be:

\[D_{out_i} = \left\lfloor\frac{D_{i} + \text{padding_before[i]} + \text{padding_after[i]} - \text{kernel_dims[i]}}{\text{stride[i]}} + 1\right\rfloor\]

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
kernel_dims (Sequence[int]) – The spatial shape of the pooling window. Only 2-D or 3-D kernel_dims are supported. If the input has int8 datatype, kernel_dims can only be 2-D.
stride (Sequence[int] | None) – A sequence of length \(M\) indicating the stride of pooling across each spatial dimension, where \(M\) is the number of spatial dimensions, i.e. \(M = \text{rank(input)} - 2\). Defaults to all 1.
padding (Sequence[Tuple[int]] | None) – A sequence of pairs of integers of length \(M\) indicating the zero padding to apply to the input along each spatial dimension before and after the dimension respectively, where \(M\) is the number of spatial dimensions, i.e. \(M = \text{rank(input)} - 2\). Defaults to all 0.

Returns:

[dtype=T1] The result tensor after the pooling operation.

Return type:

Tensor

TYPE CONSTRAINTS:

T1: float16, float32, int8

Example

Example¶

input = tp.reshape(tp.arange(16, dtype=tp.float32), (1, 1, 4, 4))
output = tp.avgpool(input, kernel_dims=(2, 2))

>>> input
tensor(
    [[[[0.0000, 1.0000, 2.0000, 3.0000],
       [4.0000, 5.0000, 6.0000, 7.0000],
       [8.0000, 9.0000, 10.0000, 11.0000],
       [12.0000, 13.0000, 14.0000, 15.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(1, 1, 4, 4))
>>> output
tensor(
    [[[[2.5000, 3.5000, 4.5000],
       [6.5000, 7.5000, 8.5000],
       [10.5000, 11.5000, 12.5000]]]], 
    dtype=float32, loc=gpu:0, shape=(1, 1, 3, 3))