ttnn.global_avg_pool2d
- ttnn.global_avg_pool2d(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig | None = None, dtype: ttnn.DataType | None = None) ttnn.Tensor
-
Applies global_avg_pool2d to
input_tensorby 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.\[\begin{split}global_avg_pool2d(\\mathrm{input\\_tensor}_i)\end{split}\]- 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
>>> tensor = ttnn.from_torch(torch.randn((10, 3, 32, 32), dtype=ttnn.bfloat16), device=device) >>> output = ttnn.global_avg_pool2d(tensor)