Pytorch NCHW/NHWC 的理解

PyTorch中的NCHW（通道-高度-宽度）和NHWC（高度-宽度-通道）是描述张量存储布局的两种方式，主要用于表示多维张量在内存中的存储顺序。理解它们对于深度学习任务中的张量操作和性能优化非常重要。

NCHW 和 NHWC 的含义

NCHW:

N: Batch size（批大小），表示在一个张量中处理的样本数量。
C: Channels（通道数），表示特征图中的通道数，例如RGB图像为3。
H: Height（高度），表示特征图的高度。
W: Width（宽度），表示特征图的宽度。

在NCHW布局中，数据的通道维度是紧跟在样本维度后面的，每个通道的像素值连续存储。

NHWC:

N: Batch size（批大小），表示在一个张量中处理的样本数量。
H: Height（高度），表示特征图的高度。
W: Width（宽度），表示特征图的宽度。
C: Channels（通道数），表示特征图中的通道数。

在NHWC布局中，每个像素的所有通道值是连续存储的，通道维度是最内层的。

理解和选择

计算效率：
NCHW 在许多GPU架构上效率更高，因为可以更有效地利用并行计算单元（如CUDA核心），特别是对于卷积操作。
NHWC 在某些情况下可能更适合CPU上的计算，因为一些CPU指令集更适合于按通道处理数据。

模型和库的兼容性：
某些深度学习框架和库可能更喜欢或更高效地使用特定的布局，因此在选择时需要考虑到模型的部署环境和所使用的库。
在PyTorch中的应用
在PyTorch中，默认情况下张量是以NCHW的格式进行存储的。可以通过转置操作（.permute()）来在NCHW和NHWC之间进行转换。例如：

import torch

# 创建一个4维张量，尺寸为(N, C, H, W)
x = torch.randn(32, 3, 224, 224)

# 转换为NHWC格式
x_nhwc = x.permute(0, 2, 3, 1)

# 再转换回NCHW格式
x_nchw = x_nhwc.permute(0, 3, 1, 2)

总结

选择合适的张量存储布局对于深度学习任务的效率和性能至关重要。NCHW和NHWC的选择取决于所使用的硬件和库的要求，同时也可以通过PyTorch提供的张量操作来进行灵活的转换和适配。