跳转到主要内容
博客

PyTorch 1.7 发布,支持 CUDA 11、FFT 新 API、Windows 分布式训练等

作者: 2020 年 10 月 27 日2024 年 11 月 16 日暂无评论

今天,我们宣布 PyTorch 1.7 以及更新的领域库已可用。PyTorch 1.7 版本包含许多新 API,包括对 NumPy 兼容的 FFT 操作、分析工具的支持,以及对分布式数据并行 (DDP) 和基于远程过程调用 (RPC) 的分布式训练的重大更新。此外,一些功能已进入稳定版,包括自定义 C++ 类、内存分析器、通过自定义类张量对象的扩展、RPC 中的用户异步函数,以及 torch.distributed 中的许多其他功能,例如每个 RPC 超时、DDP 动态分桶和 RRef 助手。

一些亮点包括:

  • CUDA 11 现已正式支持,二进制文件可在PyTorch.org 获取
  • RPC、TorchScript 和 autograd 分析器中的堆栈跟踪的分析和性能的更新和新增功能
  • (测试版) 通过 torch.fft 支持 NumPy 兼容的快速傅里叶变换 (FFT)
  • (原型) 支持 Nvidia A100 代 GPU 和原生 TF32 格式
  • (原型) Windows 上的分布式训练现已支持
  • torchvision
    • (稳定版) 变换现在支持 Tensor 输入、批量计算、GPU 和 TorchScript
    • (稳定版) JPEG 和 PNG 格式的本机图像 I/O
    • (测试版) 新的视频读取器 API
  • torchaudio
    • (稳定版) 新增语音识别 (wav2letter)、文本转语音 (WaveRNN) 和音源分离 (ConvTasNet) 支持

重申一下,从 PyTorch 1.6 开始,功能现在分为稳定版、测试版和原型版。您可以在此处查看详细公告。请注意,本博客中列出的原型功能作为此版本的一部分提供。

在此处查找完整的发行说明:此处

前端 API

[测试版] NumPy 兼容的 torch.fft 模块

FFT 相关功能常用于信号处理等各种科学领域。虽然 PyTorch 历史上支持一些 FFT 相关函数,但 1.7 版本新增了一个 torch.fft 模块,该模块实现了与 NumPy 相同 API 的 FFT 相关函数。

此新模块必须在 1.7 版本中导入才能使用,因为其名称与历史 (现已弃用) 的 torch.fft 函数冲突。

示例用法

>>> import torch.fft
>>> t = torch.arange(4)
>>> t
tensor([0, 1, 2, 3])

>>> torch.fft.fft(t)
tensor([ 6.+0.j, -2.+2.j, -2.+0.j, -2.-2.j])

>>> t = tensor([0.+1.j, 2.+3.j, 4.+5.j, 6.+7.j])
>>> torch.fft.fft(t)
tensor([12.+16.j, -8.+0.j, -4.-4.j,  0.-8.j])

[测试版] Transformer NN 模块的 C++ 支持

PyTorch 1.5以来,我们一直保持 Python 和 C++ 前端 API 的对等性。此更新允许开发人员使用 C++ 前端中的 nn.transformer 模块抽象。此外,开发人员不再需要从 Python/JIT 保存模块并加载到 C++ 中,因为现在可以直接在 C++ 中使用它。

[测试版] torch.set_deterministic

可重现性 (位对位确定性) 可能有助于在调试或测试程序时发现错误。为了促进可重现性,PyTorch 1.7 添加了torch.set_deterministic(bool)函数,该函数可以指示 PyTorch 操作员在可用时选择确定性算法,如果操作可能导致非确定性行为,则抛出运行时错误。默认情况下,此函数控制的标志为假,行为没有变化,这意味着 PyTorch 默认可能以非确定性方式实现其操作。

更准确地说,当此标志为真时:

  • 已知没有确定性实现的操作会抛出运行时错误;
  • 具有确定性变体的操作使用这些变体(通常与非确定性版本相比会降低性能);并且
  • torch.backends.cudnn.deterministic = True已设置。

请注意,这对于在单个 PyTorch 程序运行中实现确定性是必要的,但不是充分的。其他随机源,如随机数生成器、未知操作或异步或分布式计算,仍可能导致非确定性行为。

有关受影响操作的列表,请参阅torch.set_deterministic(bool)的文档。

性能与分析

[测试版] 分析器中新增堆栈跟踪

用户现在不仅可以在分析器输出表中看到操作名称/输入,还可以看到操作在代码中的位置。利用此功能的工作流几乎不需要更改。用户像以前一样使用autograd 分析器,但带有可选的新参数:with_stackgroup_by_stack_n。注意:常规分析运行不应使用此功能,因为它会增加显着的开销。

分布式训练与 RPC

[稳定版] TorchElastic 现已捆绑到 PyTorch docker 镜像中

TorchElastic 提供了当前torch.distributed.launch CLI 的严格超集,并增加了容错和弹性功能。如果用户不关心容错,他们可以通过设置max_restarts=0获得完全相同的功能/行为对等性,并增加了自动分配RANKMASTER_ADDR|PORT的便利性(而不是在torch.distributed.launch中手动指定)。

通过将torchelastic与 PyTorch 捆绑在同一个 docker 镜像中,用户无需单独安装torchelastic即可立即开始试验 TorchElastic。除了便利性之外,在现有 Kubeflow 的分布式 PyTorch 操作员中添加弹性参数支持时,这项工作也是一个不错的选择。

[测试版] DDP 中支持不均匀数据集输入

PyTorch 1.7 引入了一个新的上下文管理器,与使用torch.nn.parallel.DistributedDataParallel训练的模型结合使用,以支持在不同进程之间数据集大小不均匀的训练。此功能在使用 DDP 时提供了更大的灵活性,并防止用户必须手动确保不同进程之间的数据集大小相同。使用此上下文管理器,DDP 将自动处理不均匀的数据集大小,这可以防止训练结束时出现错误或挂起。

[测试版] NCCL 可靠性 – 异步错误/超时处理

过去,NCCL 训练运行会因为集体操作卡住而无限期挂起,给用户带来了非常不愉快的体验。如果检测到潜在的挂起,此功能将中止卡住的集体操作并抛出异常/崩溃进程。当与 torchelastic (可以从最后一个检查点恢复训练过程) 等工具一起使用时,用户可以获得更高的分布式训练可靠性。此功能是完全可选的,并且位于需要明确设置才能启用此功能的环境变量之后 (否则用户将看到与以前相同的行为)。

[测试版] TorchScriptrpc_remoterpc_sync

torch.distributed.rpc.rpc_async在之前的版本中已在 TorchScript 中可用。对于 PyTorch 1.7,此功能将扩展到其余两个核心 RPC API,即torch.distributed.rpc.rpc_synctorch.distributed.rpc.remote。这将完成 TorchScript 中主要 RPC API 的支持目标,它允许用户在 TorchScript 中使用现有的 Python RPC API(在脚本函数或脚本方法中,这会释放 Python 全局解释器锁),并可能提高多线程环境中的应用程序性能。

[测试版] 支持 TorchScript 的分布式优化器

PyTorch 为训练算法提供了广泛的优化器集,这些优化器已作为 Python API 的一部分反复使用。然而,用户通常希望使用多线程训练而不是多进程训练,因为它在大规模分布式训练(例如分布式模型并行)或任何基于 RPC 的训练应用程序中提供了更好的资源利用率和效率)。以前用户无法使用分布式优化器做到这一点,因为我们需要摆脱 Python 全局解释器锁(GIL)的限制来实现这一点。

在 PyTorch 1.7 中,我们正在分布式优化器中启用 TorchScript 支持以消除 GIL,并使在多线程应用程序中运行优化器成为可能。新的分布式优化器具有与以前完全相同的接口,但它会自动将每个工作器中的优化器转换为 TorchScript 以使每个 GIL 自由。这是通过利用函数式优化器概念并允许分布式优化器将优化器的计算部分转换为 TorchScript 来完成的。这将有助于分布式模型并行训练等用例,并提高多线程的性能。

目前,唯一支持 TorchScript 自动转换的优化器是Adagrad,所有其他优化器仍将像以前一样工作,不支持 TorchScript。我们正在努力将覆盖范围扩展到所有 PyTorch 优化器,并预计在未来版本中会有更多。启用 TorchScript 支持的使用是自动的,并且与现有的 Python API 完全相同,以下是一个如何使用它的示例

import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc
from torch import optim
from torch.distributed.optim import DistributedOptimizer

with dist_autograd.context() as context_id:
  # Forward pass.
  rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
  rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
  loss = rref1.to_here() + rref2.to_here()

  # Backward pass.
  dist_autograd.backward(context_id, [loss.sum()])

  # Optimizer, pass in optim.Adagrad, DistributedOptimizer will
  # automatically convert/compile it to TorchScript (GIL-free)
  dist_optim = DistributedOptimizer(
     optim.Adagrad,
     [rref1, rref2],
     lr=0.05,
  )
  dist_optim.step(context_id)

[测试版] 基于 RPC 的分析增强功能

PyTorch 1.6 中首次引入了结合 RPC 框架使用 PyTorch 分析器的支持。在 PyTorch 1.7 中,进行了以下增强:

  • 实现了对通过 RPC 分析 TorchScript 函数的更好支持
  • 在与 RPC 配合使用的分析器功能方面实现了对等
  • 增加了对服务器端异步 RPC 函数 (使用rpc.functions.async_execution装饰的函数) 的支持。

用户现在可以使用熟悉的分析工具,例如torch.autograd.profiler.profile()with torch.autograd.profiler.record_function,并且它与 RPC 框架透明地工作,具有完整的功能支持,分析异步函数和 TorchScript 函数。

[原型] Windows 支持分布式训练

PyTorch 1.7 带来了对 Windows 平台上DistributedDataParallel和集体通信的原型支持。在此版本中,该支持仅涵盖基于 Gloo 的ProcessGroupFileStore

要在多台机器上使用此功能,请在init_process_group中提供来自共享文件系统的一个文件。

# initialize the process group
dist.init_process_group(
    "gloo",
    # multi-machine example:
    # init_method = "file://////{machine}/{share_folder}/file"
    init_method="file:///{your local file path}",
    rank=rank,
    world_size=world_size
)

model = DistributedDataParallel(local_model, device_ids=[rank])

移动端

PyTorch Mobile 同时支持iOSAndroid,二进制包分别在CocoapodsJCenter提供。您可以在此处了解有关 PyTorch Mobile 的更多信息。

[测试版] PyTorch Mobile 缓存分配器以提高性能

在某些移动平台(例如 Pixel)上,我们观察到内存更积极地返回给系统。这导致频繁的页面错误,因为 PyTorch 作为功能框架不维护操作符的状态。因此,对于大多数操作,输出在每次执行操作时动态分配。为了改善由此带来的性能损失,PyTorch 1.7 为 CPU 提供了一个简单的缓存分配器。该分配器按张量大小缓存分配,目前仅通过 PyTorch C++ API 提供。缓存分配器本身由客户端拥有,因此分配器的生命周期也由客户端代码维护。然后,这样的客户端拥有的缓存分配器可以与作用域保护c10::WithCPUCachingAllocatorGuard一起使用,以在该作用域内启用缓存分配的使用。示例用法:

#include <c10/mobile/CPUCachingAllocator.h>
.....
c10::CPUCachingAllocator caching_allocator;
  // Owned by client code. Can be a member of some client class so as to tie the
  // the lifetime of caching allocator to that of the class.
.....
{
  c10::optional<c10::WithCPUCachingAllocatorGuard> caching_allocator_guard;
  if (FLAGS_use_caching_allocator) {
    caching_allocator_guard.emplace(&caching_allocator);
  }
  ....
  model.forward(..);
}
...

注意:缓存分配器仅在移动构建中可用,因此在移动构建之外使用缓存分配器将无效。

torchvision

[稳定版] 变换现在支持 Tensor 输入、批量计算、GPU 和 TorchScript

torchvision 变换现在继承自nn.Module,并且可以进行 torchscript 化,并应用于 torch Tensor 输入以及 PIL 图像。它们还支持带有批次维度的 Tensor,并在 CPU/GPU 设备上无缝工作

import torch
import torchvision.transforms as T

# to fix random seed, use torch.manual_seed
# instead of random.seed
torch.manual_seed(12)

transforms = torch.nn.Sequential(
    T.RandomCrop(224),
    T.RandomHorizontalFlip(p=0.3),
    T.ConvertImageDtype(torch.float),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
)
scripted_transforms = torch.jit.script(transforms)
# Note: we can similarly use T.Compose to define transforms
# transforms = T.Compose([...]) and 
# scripted_transforms = torch.jit.script(torch.nn.Sequential(*transforms.transforms))

tensor_image = torch.randint(0, 256, size=(3, 256, 256), dtype=torch.uint8)
# works directly on Tensors
out_image1 = transforms(tensor_image)
# on the GPU
out_image1_cuda = transforms(tensor_image.cuda())
# with batches
batched_image = torch.randint(0, 256, size=(4, 3, 256, 256), dtype=torch.uint8)
out_image_batched = transforms(batched_image)
# and has torchscript support
out_image2 = scripted_transforms(tensor_image)

这些改进启用了以下新功能:

  • 支持 GPU 加速
  • 批量变换,例如视频所需的变换
  • 变换多通道 torch tensor 图像(超过 3-4 个通道)
  • 将 torchscript 变换与模型一起用于部署 注意: TorchScript 支持的例外包括 ComposeRandomChoiceRandomOrderLambda 以及应用于 PIL 图像的变换,例如 ToPILImage

[稳定版] JPEG 和 PNG 格式的本地图像 IO

torchvision 0.8.0 引入了 JPEG 和 PNG 格式的本地图像读写操作。这些操作符支持 TorchScript 并返回CxHxW格式的uint8张量,因此现在可以作为模型的一部分部署在 C++ 环境中。

from torchvision.io import read_image

# tensor_image is a CxHxW uint8 Tensor
tensor_image = read_image('path_to_image.jpeg')

# or equivalently
from torchvision.io import read_file, decode_image
# raw_data is a 1d uint8 Tensor with the raw bytes
raw_data = read_file('path_to_image.jpeg')
tensor_image = decode_image(raw_data)

# all operators are torchscriptable and can be
# serialized together with your model torchscript code
scripted_read_image = torch.jit.script(read_image)

[稳定版] RetinaNet 检测模型

此版本添加了来自密集目标检测的 Focal Loss的带有 ResNet50 主干的 RetinaNet 预训练模型。

[测试版] 新的视频读取器 API

此版本引入了一种新的视频读取抽象,它提供了对视频迭代更精细的控制。它支持图像和音频,并实现了迭代器接口,因此可以与其他 Python 库(例如 itertools)互操作。

from torchvision.io import VideoReader

# stream indicates if reading from audio or video
reader = VideoReader('path_to_video.mp4', stream='video')
# can change the stream after construction
# via reader.set_current_stream

# to read all frames in a video starting at 2 seconds
for frame in reader.seek(2):
    # frame is a dict with "data" and "pts" metadata
    print(frame["data"], frame["pts"])

# because reader is an iterator you can combine it with
# itertools
from itertools import takewhile, islice
# read 10 frames starting from 2 seconds
for frame in islice(reader.seek(2), 10):
    pass
    
# or to return all frames between 2 and 5 seconds
for frame in takewhile(lambda x: x["pts"] < 5, reader):
    pass

注意事项

  • 要使用 Video Reader API 测试版,您必须从源代码编译 torchvision 并在系统中安装 ffmpeg。
  • VideoReader API 目前作为测试版发布,其 API 可能会根据用户反馈而更改。

torchaudio

通过此版本,torchaudio 正在扩展其对模型和端到端应用程序的支持,添加了 wav2letter 训练管道以及端到端文本转语音和音源分离管道。请在github上提交问题以提供反馈。

[稳定版] 语音识别

在上一个版本中添加了用于语音识别的 wav2letter 模型的基础上,我们现在添加了一个示例 wav2letter 训练管道,使用了 LibriSpeech 数据集。

[稳定版] 文本转语音

为了支持文本转语音应用程序,我们基于此存储库的实现,添加了一个基于 WaveRNN 模型的声码器。原始实现是在“高效神经音频合成”中引入的。我们还提供了一个示例 WaveRNN 训练管道,该管道使用此版本中添加到 torchaudio 的 LibriTTS 数据集。

[稳定版] 音源分离

随着 ConvTasNet 模型的添加,基于“Conv-TasNet: Surpassing Ideal Time-Frequency Magnitude Masking for Speech Separation”论文,torchaudio 现在也支持音源分离。提供了使用 wsj-mix 数据集的示例 ConvTasNet 训练管道

干杯!

PyTorch 团队