使用 Torch.Compile 静态图
最近 PyTorch 2.0 公告显示,PyTorch 团队引入了新的 toch.compile 函数。该函数可以通过生成优化的静态图来加速 PyTorch 代码执行,而不是使用动态图运行 PyTorch 代码。
由于 PyTorch 2.0 尚未正式发布,因而必须先要安装 torchtriton,并更新到 PyTorch 最新版本才能使用此功能。
然后通过添加这一行对代码进行修改:
在 4 块 GPU 上进行分布式数据并行
上文介绍了在单 GPU 上加速代码的混合精度训练,接下来介绍多 GPU 训练策略。下图总结了几种不同的多 GPU 训练技术。
想要实现分布式数据并行,可以通过 DistributedDataParallel 来实现,只需修改一行代码就能使用 Trainer。
经过这一步优化,在 4 个 A100 GPU 上,这段代码运行了 3.52 分钟就达到了 93.1% 的测试准确率。
DeepSpeed
最后,作者探索了在 Trainer 中使用深度学习优化库 DeepSpeed 以及多 GPU 策略的结果。首先必须安装 DeepSpeed 库:
接着只需更改一行代码即可启用该库:
这一波下来,用时 3.15 分钟就达到了 92.6% 的测试准确率。不过 PyTorch 也有 DeepSpeed 的替代方案:fully-sharded DataParallel,通过 strategy="fsdp" 调用,最后花费 3.62 分钟完成。
以上就是作者提高 PyTorch 模型训练速度的方法,感兴趣的小伙伴可以跟着原博客尝试一下,相信你会得到想要的结果。
原文链接:https://sebastianraschka.com/blog/2023/pytorch-faster.html