跳转至

Pytorch 5 ： Distributed Training & Parallelism

导言
  • 多GPU使用
  • 分布式训练：了解如何使用PyTorch进行多GPU训练或分布式训练。
  • 并行策略实现

在深度学习中，内存格式（memory_format）决定了张量数据在内存中的存储顺序，直接影响计算效率和硬件兼容性。以下是不同内存格式的区别及转换的作用：

1. 内存格式类型及区别

(1) contiguous_format

（默认连续格式，如 NCHW）**

  • 存储顺序：数据按维度顺序连续存储。
    • 示例：对于形状为 (N, C, H, W) 的张量（批次、通道、高度、宽度），内存中按 N→C→H→W 顺序排列。
    • 内存布局：NCHW → [N][C][H][W]。
  • 适用场景：
    • 大多数深度学习框架（如PyTorch）的默认格式。
    • 昇腾NPU、部分CPU场景下性能最佳。
  • 优点：
    • 内存连续，访问效率高。
    • 兼容性强，支持所有硬件。

(2) channels_last

（通道最后格式，如 NHWC）

  • 存储顺序：通道维度放在最后。
    • 示例：形状为 (N, H, W, C)，内存中按 N→H→W→C 顺序排列。
    • 内存布局：NHWC → [N][H][W][C]。
  • 适用场景：
    • GPU上使用Tensor Core加速时（如混合精度训练）。
    • 某些卷积操作（如Depthwise Conv）在NHWC下效率更高。
  • 优点：
    • 更适合并行计算，减少内存访问跳跃。
    • 在GPU上通常比NCHW快10%~30%。

(3) preserve_format

（保持原格式）

  • 作用：保留输入张量的现有内存格式（不主动修改）。
    • 例如：若输入是NHWC，输出仍为NHWC；若输入是NCHW，输出仍为NCHW。
  • 适用场景：
    • 当需要保持与输入一致的内存格式时（如模型中间层）。

2. 内存格式转换的作用

(1) 优化计算效率

  • 硬件适配：
  • GPU：NHWC格式更适合利用Tensor Core加速（尤其是FP16计算）。
  • NPU/ASIC：可能仅支持NCHW（如昇腾NPU），需强制使用连续格式。
  • 示例： 
    # 在GPU上使用NHWC加速卷积
x = x.to(memory_format=torch.channels_last)  # 转换为NHWC
conv_output = conv(x)  # 速度更快

(2) 减少内存碎片

  • 连续性保证：
  • contiguous() 或 to(memory_format=...) 可确保张量在内存中连续，避免因维度操作（如permute）导致内存不连续，从而减少计算错误或性能下降。

(3) 兼容性处理

  • 昇腾NPU限制：
  • 昇腾NPU的torch_npu库目前仅支持contiguous_format和preserve_format，强行使用channels_last会报错。
  • 需通过其他方式优化性能（如调整数据预处理或使用NPU专用算子）。

3. 昇腾NPU的特殊性

(1) 为何不支持 channels_last？

  • 硬件设计差异：
  • NPU的计算单元和内存控制器针对NCHW格式优化，强行使用NHWC可能导致内存访问冲突或计算效率下降。
  • 软件栈限制：
  • 昇腾AI软件栈（CANN）可能未实现NHWC格式的底层算子支持。

(2) 替代优化方案

  1. 使用默认连续格式： 
    x = x.permute(0,3,1,2).contiguous()  # 确保NCHW且内存连续
  2. 调整模型结构：
  3. 在数据输入前转换为NCHW，避免运行时转换。
  4. 使用NPU专用优化：
  5. 调用昇腾提供的高性能算子（如torch_npu.npu_format_cast）。

4. 代码示例对比

(1) GPU优化

（使用channels_last） 

# 在GPU上启用NHWC加速
x = torch.randn(1, 3, 224, 224).to("cuda").to(memory_format=torch.channels_last)
conv = nn.Conv2d(3, 64, kernel_size=3).to("cuda").to(memory_format=torch.channels_last)
output = conv(x)  # 速度更快

(2) 昇腾NPU适配

（强制连续格式） 

# 在NPU上使用默认NCHW格式
x = torch.randn(1, 3, 224, 224).npu()  # 数据加载到NPU
x = x.permute(0,3,1,2).contiguous()    # 确保连续内存
conv = nn.Conv2d(3, 64, kernel_size=3).npu()
output = conv(x)  # 兼容NPU

总结

  • 区别：contiguous_format（NCHW）是通用格式，channels_last（NHWC）适合GPU加速，preserve_format保持原格式。
  • 作用：转换内存格式可优化计算效率，但需适配硬件限制。
  • 昇腾NPU：强制使用NCHW，需通过contiguous()或调整数据预处理保证兼容性。

怎么使用GPU，怎么多GPU

在GPU上训练 就像你怎么把一个张量转移到GPU上一样，你要将神经网络转到GPU上。 如果CUDA可以用，让我们首先定义下我们的设备为第一个可见的cuda设备。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Assume that we are on a CUDA machine, then this should print a CUDA device:

print(device) # cuda:0

net=Net()
net.to(device)
outputs = net(inputs)

input = torch.randn(1, 1, 32, 32)
inputs, labels = inputs.to(device), labels.to(device)
out = net(input)

多GPU

如果你想要来看到大规模加速，使用你的所有GPU，请查看：数据并行性（https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html）。PyTorch 60 分钟入门教程：数据并行处理

http://pytorchchina.com/2018/12/11/optional-data-parallelism/