1月16日,Facebook发布了PyTorch 1.4,对音频、视觉和文本库进行了升级。 在最新版本中,PyTorch 最大的变化在于增加了支持分布式模型并行训练、为 PyTorch Mobile 提供 Build 级别的支持、torch.optim 更新等多项新的特性。
2019年旧金山PyTorch开发大会
支持分布式模型并行训练
1.4 版最大的亮点在于对分布式模型并行训练增加了支持。
官方表示,随着 RoBERTa 等万亿级别参数的大型模型出现,人们越来越需要模型并行训练了。因此本次版本会提供一个分布式 RPC(Romote Procedure Call)框架。
这一框架可以远程运行函数,在不复制真实数据的情况下查询远程对象。此外 PyTorch 还提供了 autograd 和优化器API,能够透明地运行后端并跨 RPC 边界更新参数。
相对应的,PyTorch 引入了名为 torch.distributed.rpc 的库。这是一个有基本构建单元的代码库,用于构建能够在模型训练和推断时远程运行的函数。
具体来说,这一个库有四个主要部分:RPC、Remote Reference、Distributed Autograd 和 Distributed Optimizer。
为 PyTorch Mobile 提供 Build 级别的支持
去年秋天,作为PyTorch 1.3版本的一部分,面向iOS和Android设备的PyTorch Mobile首次发布,速度提升主要由于量化、谷歌TPU支持和JIT编译器升级。
本次版本更新增加了更多对移动端的支持,如 fine-grain 级别的定制化构建脚本,这可以让移动端开发者优化代码库的大小,只包括他们的模型所使用的算子。
同时,在运行过程中显著减少对设备空间的占用。早期的结果说明,一个定制化的 MobileNetV2 比 PyTorch 的移动端库构建出来的要小 40% 到 50%。
对 Java binding 的支持
除了对 Python 和 C++的支持以外,本次更新增加了对 Java binding 的实验性支持。基于在 PyTorch Mobile 中对安卓开发的交互界面,用户可以从任何 Java 程序中调用 TorchScript 模型。
如何在Java中使用PyTorch
当然,在本次更新中,Java binding 只在 Linux 版本上存在,且只能进行模型推理,官方将会在后续更新中扩展支持。
torch.optim 更新
除了以上三个重要特性外,还有一个新的小特性值得注意。torch.optim.lr_scheduler 现已支持“链式更新”(chaining)。即用户可以定义两个 schedulers,并交替在训练中使用。
其他的新特性更新包括对分布式、JIT 等方面的更新。
在改进方面值得注意的是对 C++的多项 API 进行了更新和调整,此外还有对 JIT、ONNX 的改进,以及 AMD 硬件的支持等,另外,nn.GELU 正式成为了激活函数的一种。
C++更新的 API 包括:
- Torch::nn modules:包括卷积层、池化层、损失层、归一化层、激活层、Dropout层、嵌入层等。
- Torch::nn::functional functions:包括卷积、池化、损失、归一化、相似度等方面的函数。
在改进的过程中,PyTorch 1.4 还修复了大约 30 个 Bug,涵盖 CUDA、损失函数、卷积、嵌入等多个方面的代码。