helios.nn.layers.pool

Classes

AdaptiveAvgPool2d

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

Module Contents

class helios.nn.layers.pool.AdaptiveAvgPool2d(output_size: int | tuple[int | None, int | None])

Bases: torch.nn.Module

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

The output is of size \(H \times W\), for any input size. The number of output features is equal to the number of input planes.

Note

This class is a re-implementation of torch.nn.AdaptiveAvgPool2d that can be exported to ONNX and serves as a drop-in replacement for torch’s version.

Parameters:

output_size – the target output size of the image of the form \(H \times W\). Can be a tuple (H, W) or a single H for a square image \(H \times H\). H and W can be either a int, or None which means the size will be the same as that of the input.

forward(x: torch.Tensor) → torch.Tensor

Evaluate the adaptive average pool on the input.

Parameters:

x – the input tensor.

Returns:

The pooled tensor.