优化PyTorch性能的一些trick

1. 混合精度训练：FP32量化为FP16
2. 大Batch训练：对前几次梯度进行累加，然后统一进行参数更新，从而变相实现大Batch训练
3. 梯度检查点：训练时间换显存，在前向阶段传递到 checkpoint 中的 forward 函数会以 torch.no_grad 模式运行，并且仅仅保存输入参数和 forward 函数，在反向阶段重新计算其 forward 输出值。
4. 尽量实现向量化运算（即使开的数组多占用空间或者操作步数多），
5. 使用@torch.jit.script装饰器会将pointwise操作融合成一个kernel来执行，比如激活函数（pointwise操作）中的多个运算
6. pin_memory=True，减少数据传送时间
7. 直接在GPU上创建tensor，eg：tensor = torch.rand([10,5], device=torch.device('cuda:0'))
8. torch.tensor()始终复制数据，numpy到tensor使用torch.from_numpy(numpy_array)更快，tensor（或list、tuple）到tensor使用torch.as_tensor(others)更快
9. 使用tensor.to(device, non_blocking=True)允许异步来overlap data transfers and kernel execution，数据流向构成一个DAG
10. 输入、输出、中间参数的维度尽量是8的倍数
11. 如果输入大小不经常变化，在训练前设置torch.backends.cudnn.benchmark=True，cuDNN会自动寻找最好的算法来计算卷积（因为cuDNN算法计算size不同的卷积的性能不同）
12. 在BN层之前的卷积层中不使用bias（因为BN第一步会减去均值）
13. Use channels_last memory format for 4D NCHW Tensors
14. use DistributedDataParallel instead of DataParallel

参考资料：

推荐：Optimize PyTorch Performance for Speed and Memory Effciency(2022)

拿什么拯救我的4G显卡

用PyTorch踩过的坑

Effective PyTorch