跳转到主要内容
博客

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 和自动求导分析器中栈跟踪的性能分析更新和新增功能
  • (Beta) 通过 torch.fft 支持 NumPy 兼容的快速傅里叶变换 (FFT)
  • (原型) 支持 Nvidia A100 代 GPU 和原生 TF32 格式
  • (原型) Windows 上分布式训练现已支持
  • torchvision
    • (稳定版) Transforms 现已支持 Tensor 输入、批处理计算、GPU 和 TorchScript
    • (稳定版) JPEG 和 PNG 格式的原生图像 I/O
    • (Beta) 新的视频阅读器 API
  • torchaudio
    • (稳定版) 新增语音识别 (wav2letter)、文本转语音 (WaveRNN) 和音源分离 (ConvTasNet) 支持

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

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

前端 API

[Beta] 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])

[Beta] C++ 支持 Transformer NN 模块

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

[Beta] torch.set_deterministic

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

更准确地说，当此标志为 true 时：

  • 已知没有确定性实现的操作会抛出运行时错误；
  • 具有确定性变体的操作使用这些变体（通常与非确定性版本相比有性能损失）；以及
  • torch.backends.cudnn.deterministic = True已设置。

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

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

性能与分析

[Beta] 栈跟踪已添加到分析器

用户现在不仅可以在分析器输出表中看到运算符名称/输入，还可以看到运算符在代码中的位置。利用此功能的工作流程所需更改极少。用户像以前一样使用自动求导分析器，但带有可选的新参数：with_stack和group_by_stack_n。注意：常规分析运行不应使用此功能，因为它会增加显着开销。

分布式训练与 RPC

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

TorchElastic 提供了当前torch.distributed.launch CLI 的严格超集，并增加了容错和弹性功能。如果用户对容错不感兴趣，他们可以通过将max_restarts=0设置为获得完全相同的功能/行为，并额外提供自动分配的RANK和MASTER_ADDR|PORT（而无需在torch.distributed.launch中手动指定）。

通过将torchelastic捆绑在与 PyTorch 相同的 docker 镜像中，用户可以立即开始尝试 TorchElastic，而无需单独安装torchelastic。除了方便之外，这项工作在现有的 Kubeflow 分布式 PyTorch 运算符中添加弹性参数支持时也很有用。

[Beta] DDP 中支持不均匀数据集输入

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

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

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

[Beta] TorchScriptrpc_remote和rpc_sync

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

[Beta] 支持 TorchScript 的分布式优化器

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

在 PyTorch 1.7 中，我们正在分布式优化器中启用 TorchScript 支持以消除 GIL，并使其能够在多线程应用程序中运行优化器。新的分布式优化器具有与以前完全相同的接口，但它会自动将每个 worker 中的优化器转换为 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)

[Beta] 基于 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 的ProcessGroup和FileStore。

要在多台机器上使用此功能，请在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 支持 iOS 和 Android，二进制包分别在 Cocoapods 和 JCenter 中提供。您可以在此处了解更多关于 PyTorch Mobile 的信息。

[Beta] 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

[稳定版] Transforms 现已支持 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 张量图像（具有超过 3-4 个通道）
  • 将转换与您的模型一起 torchscript 以进行部署**注意：**TorchScript 支持的例外包括Compose、RandomChoice、RandomOrder、Lambda以及应用于 PIL 图像的转换，例如ToPILImage。

[稳定版] JPEG 和 PNG 格式的原生图像 I/O

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 预训练模型。

[Beta] 新的视频阅读器 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

注意事项

  • 为了使用视频阅读器 API beta 版，您必须从源代码编译 torchvision 并在系统中安装 ffmpeg。
  • VideoReader API 目前作为 beta 版发布，其 API 可能会根据用户反馈而改变。

torchaudio

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

[稳定版] 语音识别

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

[稳定版] 文本转语音

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

[稳定版] 音源分离

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

干杯！

PyTorch 团队