13、Pytorch 模型训练加速_pytorch模型训练性能调优宝典

---

一、优化 I/O 性能

• I/O 是训练过程中的主要瓶颈， 这是因为数据必须 从存储加载到GPU 进行处理。 特别是当数据集很大或位于远程存储时， 这个过程可能非常耗时
• 如果 I/O 速度较慢， 在数据准备好前， GPU 会处于空闲状态， 导致昂贵 GPU 资源的浪费。 我们观察到， 一般在 I/O 缓慢的情况下， GPU的利用率低于 50%
• 优化 I/O 的方式：
  • 将数据拷贝到 SSD，单个NVMe（ SSD） 硬盘可以提供高达 7 Gbps的输入/输出（ I/O） 带宽
  • 启用异步数据加载，将 DataLoder 中的 num_workers 的值设置大于 0，且 pin_memory 设置为 True（可以通过使用固定内存来实现从主机到 GPU 的异步内存拷贝）

---

二、数据操作优化

• 强烈建议直接在打算执行操作的设备上创建张量 。例如下面的代码， 其中张量最初在 CPU 内存中创建 ， 然后被移动到 GPU 显存 。 这种方法通常需要分配 CPU RAM 并使用 cudaMemcpy函数在 GPU 和 CPU 之间传输数据。通过直接在 GPU 上创建张量， 可以避免不必要的内存数据传输， 进而提升性能并降低开销。该方法在处理深度学习任务中涉及大量GPU操作时特别有效。
  

• 使用 torch.as_tensor：与 torch.tensor() 不同， 它能避免不必要的数据拷贝， 将数据转换为张量， 并在可能的情况下共享底层数据和保留 autograd 历史记录

• 将 non_blocking 设置为 True：并行加载数据到 GPU 设备内存
  

---

三、针对 GPU 的优化

• 您可以将想要优化的任何 Python 函数传递给 torch.compile， 然后使用 torch.compile 的返回值作为优化函数来替换原始函数

• 如果不想进行显式替换， 也可以使用“@torch.compile” 装饰现有函数， 之后可以在现有代码中继续使用原始函数名
  
  

• 使用 DistributedDataParallel（ DDP）：可跨多个 GPU 或机器实现分布式训练， 从而显著地提高计算能力并缩短训练时间

• 使用低精度的数据类型
  

with autocast(device_type='cuda', dtype=torch.float16):
	output = model(input)
loss = loss_fn(output, target)

# 注意，这里不推荐在autocast区域执行反向传播， 因为反向操作应当以对应的（ autocast 所选择的） 前向传播操作的数据类型运行
1
2
3
4
5

---

四、针对 CPU 的优化

• Pillow 作为 PIL 的直接替代品， 已经因其超越其他图像处理库的速度优势而闻名， 而 PillowSIMD 可以进一步提升性能。 基准测试显示， 在相同硬件/平台上， Pillow-SIMD 比 Pillow 和其他库（ 如ImageMagick, OpenCV, and IPP） 快4-6倍（pytorch 官方的 torchvision 使用的是 pillow 来读取图片和做预处理，它默认只是给新手入门的并不适合真正的做大规模实验）。 它通过使用可同时并行处理多数据点的 SIMD 技术来优化常见的图像操作指令， 从而实现速度上的提升。
• SIMD 代表 single instruction, multiple data, 充分压榨多核（Intel） CPU 的性能. 常用的 CPU SIMD 指令集有
  MMX, SSE-SSE4, AVX, AVX2, AVX512, NEON 等。Pillow-SIMD 就是 应用 SIMD 加速的 Pillow，100% API cover， production ready。默认 SSE4 编译，提供 AVX2 支持

# 已加速常用操作（主要针对的是图片预处理， Image.open() 并没有加速）
Resize (convolution-based resampling): SSE4, AVX2
Gaussian and box blur: SSE4
Alpha composition: SSE4, AVX2
RGBA → RGBa (alpha premultiplication): SSE4, AVX2
RGBa → RGBA (division by alpha): SSE4, AVX2
RGB → L (grayscale): SSE4
3x3 and 5x5 kernel filters: SSE4, AVX2
Split and get_channel: SSE4


# 安装流程
# must remove existed pillow first.
$ pip uninstall pillow

# install SSE4 version
$ pip install pillow-simd

# install AVX2 version
$ CC="cc -mavx2" pip install -U --force-reinstall pillow-simd

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

---

五、参考资料

1、https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html
2、https://docs.nvidia.com/deeplearning/performance/index.html
3、https://github.com/uploadcare/pillow-simd
4、PyTorch中，18个速度和内存效率优化技巧