汤不热吧PyTorch 默认使用 NCHW (Batch, Channels, Height, Width) 内存布局，这是一种从科学计算历史遗留下来的传统格式。然而，对于现代视觉模型，尤其是在利用 NVIDIA Tensor Core 或其他高度优化的深度学习加速器进行推理时，NCHW 格式常常不是最优解。
本文将详细介绍如何利用 PyTorch 提供的 torch.channels_last (即 NHWC) 内存格式，来加速卷积神经网络（CNN）的运算，并解释其背后的底层原理。
什么是 Channels Last (NHWC)？
数据在内存中的存储顺序对于性能至关重要。对于一个四维张量 (N, C, H, W)：
- NCHW (默认格式): 维度顺序是 C（通道）在 H（高）和 W（宽）之前。内存中存储时，所有像素点的数据是连续的，但当你从一个像素点移动到其下一个通道的特征时，内存跳跃距离很长。
- NHWC (Channels Last 格式): 维度顺序是 H, W 在 C 之前。这意味着对于图像中的一个特定像素 (h, w)，其所有通道的特征值 C 是在内存中连续存储的。
Channels Last 加速的底层原理
NHWC 格式能够加速视觉模型卷积运算，主要基于以下两个核心原因：
1. 增强数据局部性（Data Locality）
在 NHWC 格式中，同一个空间位置（H, W）上的所有通道数据是内存连续的。当进行卷积操作时，硬件（如 GPU 缓存）可以一次性加载与该像素相关的所有通道信息。这种高局部性大大减少了 CPU 或 GPU 访问主内存的次数，提高了缓存命中率。
2. 更好地适配硬件优化（尤其是 Tensor Cores）
许多现代深度学习硬件（如 NVIDIA GPU 上的 Tensor Cores）内部将卷积操作转化为高效的矩阵乘法（GEMM）。这些硬件通常针对 NHWC 格式设计了更高效的矩阵重排和计算流程，因为 NHWC 格式的张量结构更自然地映射到这些硬件的计算单元布局，从而实现了更快的计算速度。
实践：如何将 PyTorch 模型和数据切换到 Channels Last
在 PyTorch 中启用 Channels Last 非常简单，只需要对模型和输入数据使用 .to(memory_format=torch.channels_last) 方法即可。需要注意的是，必须同时转换模型和输入张量。
下面的代码演示了如何在 PyTorch 中应用和测试 Channels Last 优化（建议在拥有 GPU 的环境中运行，效果最明显）：
import torch
import torch.nn as nn
import time
# 1. 定义一个简单的 CNN 模型
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
self.relu = nn.ReLU()
self.pool = nn.MaxPool2d(2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
def forward(self, x):
x = self.pool(self.relu(self.conv1(x)))
x = self.pool(self.relu(self.conv2(x)))
return x
# 2. 模型和数据准备 (使用GPU进行加速测试)
B, C, H, W = 64, 3, 224, 224 # 大批量/大图像尺寸，更易看出差异
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 确保在评估模式下进行测试
model = SimpleCNN().to(device).eval()
input_data = torch.randn(B, C, H, W, device=device)
# 3. 定义性能测试函数
def benchmark(model, data, iterations=50):
# 预热
for _ in range(5):
_ = model(data)
start = time.time()
for _ in range(iterations):
_ = model(data)
if device.type == 'cuda':
torch.cuda.synchronize() # 确保CUDA操作完成
end = time.time()
return (end - start) * 1000 / iterations # 平均时间 (ms)
# --- NCHW (默认格式) 测试 ---
print("\n--- 运行 NCHW 格式 ---")
nchw_time = benchmark(model, input_data)
print(f"NCHW Average Time: {nchw_time:.3f} ms")
# --- Channels Last (NHWC) 格式转换与测试 ---
# 4. 转换模型权重和输入数据到 Channels Last 格式
model_nhwc = SimpleCNN().to(device).eval()
model_nhwc = model_nhwc.to(memory_format=torch.channels_last)
# 关键：将输入数据也转为 Channels Last
input_nhwc = input_data.contiguous(memory_format=torch.channels_last)
print("\n--- 运行 Channels Last (NHWC) 格式 ---")
nhwc_time = benchmark(model_nhwc, input_nhwc)
print(f"Channels Last Average Time: {nhwc_time:.3f} ms")
# 5. 结果对比
if nhwc_time < nchw_time:
speedup = ((nchw_time - nhwc_time) / nchw_time) * 100
print(f"\nChannels Last 成功提速: {speedup:.2f}%")
else:
print("\nChannels Last 未观察到明显提速 (可能因为CPU环境或模型过小)")
总结与适用场景
Channels Last 格式对于大多数视觉任务的 CNN 模型（如 ResNet, VGG, EfficientNet 等）在 GPU 上进行推理时，能提供显著的性能提升。它通过优化内存布局，提高了数据局部性，并更好地利用了现代硬件的矩阵乘法能力。