跳转到主要内容
博客

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

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

今天,我们宣布 PyTorch 1.6 版本正式发布,同时更新了领域库。我们也很高兴地宣布 Microsoft 团队现已负责维护 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 Distributed 支持两种强大的范例:用于模型全同步数据并行训练的 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 版本带来了对复数张量的 Beta 级支持,包括 torch.complex64 和 torch.complex128 dtypes。复数是一种可以表示为 a + bj 形式的数字,其中 a 和 b 是实数,j 是方程 x^2 = −1 的解。复数在数学和工程中经常出现,尤其是在信号处理中,而复数神经网络是一个活跃的研究领域。复数张量的 Beta 版本将支持常见的 PyTorch 和复数张量功能,以及 Torchaudio、ESPnet 等所需的函数。虽然这是此功能的早期版本,我们预计它会随着时间的推移而改进,但总体目标是提供与 NumPy 兼容的用户体验,该体验利用 PyTorch 在加速器上运行并与自动求导协作的能力,以更好地支持科学社区。

移动端更新

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

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

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

更新的领域库

torchvision 0.7

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

torchaudio 0.6

torchaudio 现已正式支持 Windows。此版本还引入了一个新的模型模块(包含 wav2letter)、新的函数(contrast、cvm、dcshift、overdrive、vad、phaser、flanger、biquad)、数据集(GTZAN、CMU)以及一个新的可选 sox 后端,支持 TorchScript。

附加更新

黑客马拉松

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

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

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

LPCV 挑战赛

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

原型功能

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

分布式 RPC/分析器

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

TorchScript 模块冻结

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

图模式量化

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

量化数值套件

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

干杯!

PyTorch 团队