PyTorch 1.9强化分布式训练和科学运算支持

由脸书发起的开源深度学习框架PyTorch,现在发布最新1.9版本,自1.8版本以来,由398位贡献者发出3,400次程序代码提交完成,该框架发展迅速,这次的更新内容也是相当丰富,包括大幅改进对科学运算的支持外,还加入新的行动解释器,来改善移动设备上的二元档大小,并且强化GPU大规模分布式训练的支持。

这次更新的科学运算前端API,包括稳定版本的torch.linalg、Complex Autograd和测试版的torch.special。由于线性代数是深度学习和科学运算中,不可或缺的工具,因此官方加入了torch.linalg模块,来扩展了PyTorch对线性代数的支持,该模块实例了NumPy中线性代数的所有函数,因此对于原本就熟悉NumPy的开发者,可以快速上手。

Complex Autograd在PyTorch 1.8的时候,就以测试版本加入,现在已经进入稳定阶段,自测试版以来,官方已经为98%的运算符加入Complex Autograd的支持。Complex Autograd能自动计算复杂梯度,并且优化实值(Real Valued)损失函数,能够对PyTorch下游项目包括TorchAudio、ESPNet、Asteroid和FastMRI,提供更好的运算支持。还在测试阶段的torch.special,则类似SciPy特殊模块,该模块包含许多科学运算和好用的函数。

移动设备部分的改进,最大的亮点在于加入了测试版行动解释器,这是PyTorch Runtime的轻量版本,可在边缘设备执行PyTorch程序,并且大幅缩减二元文件大小。官方提到,行动解释器是众所期待的功能,目前在Android中使用MobileNetV2的pt格式文件,在未压缩前17.8 MB,压缩后是8.6 MB,而通过使用行动解释器,未压缩文件大小会是8 MB,压缩后会落在4 MB左右。

在分布式训练组件方面,TorchElastic终于成为PyTorch核心的一部分,TorchElastic已经在GitHub pytorch/elastic存储库中开源一年多,现在终于正式进入PyTorch核心。TorchElastic的核心功能便是能够有弹性地处理缩放事件,可以说是PyTorch工作程序的执行器和协调器,并且能让用户使用先占式执行实例,进行分布式训练。

除了TorchElastic,分布式训练更新还包括在RPC(Remote Procedure Call)中提供CUDA支持,官方提到,与CPU RPC和通用RPC框架相比,CUDA RPC是一个更具效率的P2P张量通信方法,甚至可实现CPU RPC的34倍速。CUDA RPC创建在TensorPipe上,可以根据Tensor设备类型和信道可用性,自动替每个张量选择通信信道。

PyTorch 1.9加入了新的模型打包方法,开发者使用torch.package就能够以自包含且稳定的格式,打包PyTorch模型,模型包中会包含参数等模型资料,以及模型架构程序代码。官方提到,用torch.package可以完整打包模型的相依项目,并且加入带有Conda环境的描述档,能简单地重现训练工作负载。