torch pooling does two NHWC<->NCHW copies when one would do
On the torch backend, the pooling ops in keras/src/backend/torch/nn.py (max_pool, average_pool, adaptive_max_pool, adaptive_average_pool) call _transpose_spatial_inputs(inputs) in the default channels_last data format, which permutes NHWC->NCHW and then calls a plain .contiguous().
That .contiguous() is a full physical copy on every call. It is also unnecessary: a permuted NHWC->NCHW view of a contiguous input already has channels_last strides, so requesting torch.channels_last instead copies nothing. On top of that, the PyTorch ATen 2D pooling kernels (max_pool2d, avg_pool2d, adaptive_max_pool2d, adaptive_avg_pool2d) propagate channels_last to their output, so _transpose_spatial_outputs then returns an already-contiguous NHWC tensor rather than a non-contiguous permuted view.
This is exactly the same redundancy that #23279 identified for conv(), and it is not data-dependent: it fires on every pooling forward pass in channels_last.
_transpose_spatial_inputs already grew a channels_last_memory_format parameter in #23199 for the conv path; the pooling ops simply never opted in.
Measured on torch 2.11.0, input (32, 112, 112, 64), data_format="channels_last", outputs bit-exact (max diff = 0.0):
| op | CPU before | CPU after | MPS before | MPS after |
|---|---|---|---|---|
max_pool (valid) |
9.37 ms | 2.01 ms | 2.21 ms | 1.68 ms |
average_pool (valid) |
7.59 ms | 1.18 ms | 2.22 ms | 1.11 ms |
adaptive_max_pool |
7.07 ms | 2.03 ms | 2.24 ms | 1.63 ms |
adaptive_average_pool |
5.78 ms | 1.18 ms | 2.20 ms | 1.11 ms |
Fix: pass channels_last_memory_format=True to _transpose_spatial_inputs in the four pooling ops, matching what #23199 did for conv().
Part of the per-call Python-dispatch overhead series in #22561.
Source: keras-team/keras