#23649·keras

torch pooling 在一个卷积层中执行两次 NHWC<->NCHW 拷贝,而单个卷积层只需执行一次

作者: pctablet505创建于 2026年9月16日更新于 2026年9月16日
标签backend:torchlayers

在 Torch 后端中,`keras/src/backend/torch/nn.py` 中的池化操作(`max_pool`,`average_pool`,`adaptive_max_pool`,`adaptive_average_pool`)调用了默认 `channels_last` 数据格式中的 `_transpose_spatial_inputs(inputs)` 函数,该函数会将 NHWC->NCHW 进行转置,然后调用普通的 `.contiguous()` 函数。该 `.contiguous()` 函数在每次调用中都会进行完全的物理复制。此外,这也是不必要的: 一个连续输入的转置后的 NHWC->NCHW 视图已经具有 `channels_last` 步长,因此请求 `torch.channels_last` 并不会进行任何复制。此外,PyTorch ATen 2D 池化内核(`max_pool2d`,`avg_pool2d`,`adaptive_max_pool2d`,`adaptive_avg_pool2d`)会将 `channels_last` 传递到其输出,因此 `_transpose_spatial_outputs` 将返回已经连续的 NHWC 向量,而不是非连续的转置视图。

内容来源: keras-team/keras