ttnn.global_avg_pool2d

ttnn.global_avg_pool2d(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = None, dtype: ttnn.DataType = None) → ttnn.Tensor

Applies global_avg_pool2d to input_tensor by performing a 2D adaptive average pooling over an input signal composed of several input planes. This operation computes the average of all elements in each channel across the entire spatial dimensions.

\[global\_avg\_pool(\mathrm{input\_tensor}_i)\]
Parameters:

input_tensor (ttnn.Tensor) – the input tensor. Typically of shape (batch_size, channels, height, width).

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the operation. Defaults to None.

  • dtype (ttnn.DataType, optional) – data type for the output tensor. Defaults to None

Returns:

ttnn.Tensor – the output tensor with the averaged values. The output tensor shape is (batch_size, channels, 1, 1).

Example

# Create a random input tensor
tensor = ttnn.rand((10, 3, 32, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)

# Perform global average pooling
output = ttnn.global_avg_pool2d(tensor)
logger.info(f"Output: {output}")