1月16日,Facebook发布了PyTorch 1.4,对音频、视觉和文本库进行了升级。 在最新版本中,PyTorch 最大的变化在于增加了支持分布式模型并行训练、为 PyTorch Mobile 提供 Build 级别的支持、torch.optim 更新等多项新的特性。




Pytorch 手机端模型部署 pytorch移动端模型部署_Java

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 模型。




Pytorch 手机端模型部署 pytorch移动端模型部署_Pytorch 手机端模型部署_02

如何在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、损失函数、卷积、嵌入等多个方面的代码。