跳转到主要内容
博客

PyTorch 1.6 发布,支持原生 AMP,微软加入 Windows 维护者行列

作者: 2020 年 7 月 28 日2024 年 11 月 16 日暂无评论

今天,我们宣布 PyTorch 1.6 发布,并更新了领域库。我们也很高兴地宣布,微软团队现在负责维护 Windows 版本和二进制文件,并将通过 GitHub 和 PyTorch Windows 讨论论坛为社区提供支持。

PyTorch 1.6 版本包含许多新的 API、性能改进和分析工具,以及分布式数据并行 (DDP) 和基于远程过程调用 (RPC) 的分布式训练的重大更新。其中一些亮点包括:

  1. 自动混合精度 (AMP) 训练现在原生支持并成为稳定功能(详见此处)——感谢 NVIDIA 的贡献;
  2. 现在增加了原生 TensorPipe 支持,用于专门为机器学习构建的张量感知、点对点通信原语;
  3. 为前端 API 表面添加了对复数张量的支持;
  4. 新的分析工具提供张量级内存消耗信息;
  5. 分布式数据并行 (DDP) 训练和远程过程调用 (RPC) 包的诸多改进和新功能。

此外,从这个版本开始,功能将分为稳定版、测试版和原型版。原型版功能不作为二进制发行版的一部分,而是通过从源代码构建、使用每日构建或通过编译器标志提供。您可以在此处的帖子中了解此更改的含义。您还可以在此处找到完整的发行说明。

性能与分析

[稳定版] 自动混合精度 (AMP) 训练

AMP 允许用户轻松启用自动混合精度训练,从而在 Tensor Core GPU 上实现更高的性能和高达 50% 的内存节省。使用原生支持的 torch.cuda.amp API,AMP 提供混合精度便捷方法,其中一些操作使用 torch.float32 (float) 数据类型,而其他操作使用 torch.float16 (half)。一些操作(如线性层和卷积)在 float16 中速度更快。其他操作(如归约)通常需要 float32 的动态范围。混合精度尝试将每个操作与其适当的数据类型匹配。

[测试版] Fork/Join 并行

此版本增加了对 TorchScript 代码中语言级构造以及运行时对粗粒度并行的支持。此支持对于并行运行模型集合或并行运行循环网络的双向组件等情况非常有用,并允许解锁并行架构(例如多核 CPU)的计算能力以实现任务级并行。

TorchScript 程序的并行执行通过两个原语启用:torch.jit.forktorch.jit.wait。在下面的示例中,我们并行执行 foo

import torch
from typing import List

def foo(x):
    return torch.neg(x)

@torch.jit.script
def example(x):
    futures = [torch.jit.fork(foo, x) for _ in range(100)]
    results = [torch.jit.wait(future) for future in futures]
    return torch.sum(torch.stack(results))

print(example(torch.ones([])))

[测试版] 内存分析器

torch.autograd.profiler API 现在包含一个内存分析器,可让您检查 CPU 和 GPU 模型中不同操作符的张量内存成本。

这是一个 API 使用示例

import torch
import torchvision.models as models
import torch.autograd.profiler as profiler

model = models.resnet18()
inputs = torch.randn(5, 3, 224, 224)
with profiler.profile(profile_memory=True, record_shapes=True) as prof:
    model(inputs)

# NOTE: some columns were removed for brevity
print(prof.key_averages().table(sort_by="self_cpu_memory_usage", row_limit=10))
# ---------------------------  ---------------  ---------------  ---------------
# Name                         CPU Mem          Self CPU Mem     Number of Calls
# ---------------------------  ---------------  ---------------  ---------------
# empty                        94.79 Mb         94.79 Mb         123
# resize_                      11.48 Mb         11.48 Mb         2
# addmm                        19.53 Kb         19.53 Kb         1
# empty_strided                4 b              4 b              1
# conv2d                       47.37 Mb         0 b              20
# ---------------------------  ---------------  ---------------  ---------------

分布式训练与 RPC

[测试版] RPC 的 TensorPipe 后端

PyTorch 1.6 为 RPC 模块引入了一个新后端,该后端利用 TensorPipe 库,这是一个面向机器学习的张量感知点对点通信原语,旨在补充 PyTorch 中用于分布式训练的现有原语(Gloo、MPI 等),这些原语是集体且阻塞的。TensorPipe 的成对和异步特性使其适用于超越数据并行的新网络范例:客户端-服务器方法(例如,用于嵌入的参数服务器,Impala 风格 RL 中的 actor-learner 分离等)以及模型和管道并行训练(想想 GPipe)、gossip SGD 等。

# One-line change needed to opt in
torch.distributed.rpc.init_rpc(
    ...
    backend=torch.distributed.rpc.BackendType.TENSORPIPE,
)

# No changes to the rest of the RPC API
torch.distributed.rpc.rpc_sync(...)

[测试版] DDP+RPC

PyTorch 分布式支持两种强大的范式:用于模型完全同步数据并行训练的 DDP 和允许分布式模型并行的 RPC 框架。以前,这两个功能是独立工作的,用户无法混合搭配它们来尝试混合并行范式。

从 PyTorch 1.6 开始,我们启用了 DDP 和 RPC 无缝协作,以便用户可以结合这两种技术来实现数据并行和模型并行。一个例子是用户希望将大型嵌入表放置在参数服务器上,并使用 RPC 框架进行嵌入查找,但将较小的密集参数存储在训练器上,并使用 DDP 同步密集参数。下面是一个简单的代码片段。

// On each trainer

remote_emb = create_emb(on="ps", ...)
ddp_model = DDP(dense_model)

for data in batch:
   with torch.distributed.autograd.context():
      res = remote_emb(data)
      loss = ddp_model(res)
      torch.distributed.autograd.backward([loss])

[测试版] RPC – 异步用户函数

RPC 异步用户函数支持在执行用户定义函数时在服务器端暂停和恢复。在此功能之前,当被调用方处理请求时,一个 RPC 线程会等待直到用户函数返回。如果用户函数包含 IO(例如嵌套 RPC)或信号(例如等待另一个请求解除阻塞),则相应的 RPC 线程将空闲等待这些事件。结果,一些应用程序不得不使用非常大量的线程并发送额外的 RPC 请求,这可能导致性能下降。要使用户函数在这些事件上暂停,应用程序需要:1) 使用 @rpc.functions.async_execution 装饰器装饰函数;和 2) 让函数返回一个 torch.futures.Future 并在 Future 对象上安装恢复逻辑作为回调。请参阅下面的示例

@rpc.functions.async_execution
def async_add_chained(to, x, y, z):
    return rpc.rpc_async(to, torch.add, args=(x, y)).then(
        lambda fut: fut.wait() + z
    )

ret = rpc.rpc_sync(
    "worker1", 
    async_add_chained, 
    args=("worker2", torch.ones(2), 1, 1)
)
        
print(ret)  # prints tensor([3., 3.])
  • 使用异步用户函数进行高性能批处理 RPC 的教程
  • 文档 (链接)
  • 使用示例 (链接)

前端 API 更新

[测试版] 复数

PyTorch 1.6 版本带来了对复数张量的测试版支持,包括 torch.complex64 和 torch.complex128 dtypes。复数是可以表示为 a + bj 形式的数字,其中 a 和 b 是实数,j 是方程 x^2 = -1 的解。复数在数学和工程中经常出现,特别是在信号处理领域,而复数神经网络是一个活跃的研究领域。复数张量的测试版将支持常见的 PyTorch 和复数张量功能,以及 Torchaudio、ESPnet 等所需的函数。虽然这是此功能的早期版本,我们预计它会随着时间的推移而改进,但总体目标是提供与 NumPy 兼容的用户体验,该体验利用 PyTorch 在加速器上运行和与 autograd 协作的能力,以更好地支持科学社区。

移动端更新

PyTorch 1.6 提高了移动端设备推理的性能和整体稳定性。我们修复了一些错误,继续维护并添加了一些新功能,同时改进了在 CPU 后端各种 ML 模型推理上的 fp32 和 int8 性能。

[测试版] 移动端功能和性能

  • 无状态和有状态 XNNPACK Conv 和 Linear 算子
  • 无状态 MaxPool2d + JIT 优化通道
  • JIT 通道优化:Conv + BatchNorm 融合、图重写以用 xnnpack 算子替换 conv2d/linear、relu/hardtanh 融合、dropout 删除
  • QNNPACK 集成消除了重新量化比例约束
  • 用于 conv、linear 和动态 linear 的每通道量化
  • 禁用移动客户端跟踪以在完整 JIT 构建上节省约 600 KB

更新的领域库

torchvision 0.7

torchvision 0.7 引入了两个新的预训练语义分割模型,FCN ResNet50DeepLabV3 ResNet50,两者均在 COCO 上训练,并且比 ResNet101 主干占用更少的内存。我们还为 torchvision 模型和算子引入了 AMP(自动混合精度)自动转换支持,该支持自动为不同的 GPU 操作选择浮点精度,以提高性能同时保持准确性。

torchaudio 0.6

torchaudio 现在正式支持 Windows。此版本还引入了一个新的模型模块(包含 wav2letter)、新的函数(对比度、cvm、dcshift、过载、vad、移相器、镶边、双二阶)、数据集(GTZAN、CMU)以及一个支持 TorchScript 的新的可选 sox 后端。

其他更新

黑客马拉松

全球 PyTorch 夏季黑客马拉松又回来了!今年,团队可以在三个类别中进行虚拟竞争

  1. PyTorch 开发工具:旨在提高研究人员和开发人员使用 PyTorch 的生产力和效率的工具或库
  2. 由 PyTorch 驱动的 Web/移动应用程序:具有 Web/移动界面和/或由 PyTorch 驱动的嵌入式设备的应用程序
  3. PyTorch 负责任 AI 开发工具:用于负责任 AI 开发的工具、库或 Web/移动应用程序

这是一个与社区建立联系并练习机器学习技能的绝佳机会。

LPCV 挑战赛

2020 CVPR 低功耗视觉挑战赛 (LPCV) – 无人机视频在线赛道的提交截止日期即将到来。您可以在 2020 年 7 月 31 日之前构建一个系统,该系统可以使用 PyTorch 和 Raspberry Pi 3B+ 准确地识别和识别无人机 (UAV) 拍摄视频中的字符。

原型功能

重申一下,PyTorch 中的原型功能是早期功能,我们正在寻求反馈,评估其有用性,并在将其升级到 Beta 或稳定版之前进行改进。以下功能不属于 PyTorch 1.6 版本,而是通过每日构建提供,并附有单独的文档/教程,以帮助促进早期使用和反馈。

分布式 RPC/分析器

允许用户使用自动梯度分析器分析使用 torch.distributed.rpc 的训练作业,并远程调用分析器以收集不同节点上的分析信息。RFC 可以在此处找到,有关如何使用此功能的简短教程可以在此处找到。

TorchScript 模块冻结

模块冻结是将模块参数和属性值内联到 TorchScript 内部表示的过程。参数和属性值被视为最终值,不能在冻结模块中修改。此功能的 PR 可以在此处找到,有关如何使用此功能的简短教程可以在此处找到。

图模式量化

Eager 模式量化要求用户更改模型,包括显式量化激活、模块融合、使用 Functional Modules 重写 torch 操作以及不支持函数量化。如果我们可以跟踪或脚本模型,那么可以使用图模式量化自动完成量化,而无需 eager 模式中的任何复杂性,并且可以通过 qconfig_dict 进行配置。有关如何使用此功能的教程可以在此处找到。

量化数值套件

量化在有效时很好,但在不满足预期精度时很难知道出了什么问题。现在提供了一个数值套件原型,用于测量量化模块和浮点模块之间的比较统计数据。这可以在 eager 模式下进行测试,并且仅在 CPU 上支持,未来将提供更多支持。有关如何使用此功能的教程可以在此处找到。

干杯!

PyTorch 团队