PyTorch 1.7 发布！ 支持CUDA 11，Windows 分布式训练，以及FFT新API

作者：Team PyTorch 编译：McGL

今天，我们正式发布 PyTorch 1.7，以及升级的域库。PyTorch 1.7版本包括了一些新的 API，比如对兼容 numpy 的 FFT 操作的支持、性能分析工具以及对分布式数据并行(DDP)和基于远程过程调用(RPC)的分布式训练的重要更新。此外，还有一些特性移到了 stable 状态，包括自定义 C++ 类、内存分析器、通过自定义类张量对象实现的扩展、 RPC 中的用户异步函数以及 torch.distributed 中的其他一些特性，如 Per-RPC 超时、 DDP dynamic bucketing 和 RRef helper。

其中的一些亮点包括:

• 正式支持 CUDA 11 ，二进制文件可在 PyTorch.org 下载
• 在 autograd 分析器中更新和添加 RPC、 TorchScript 和 Stack traces 的分析和性能
• (Beta) 通过 torch.fft 支持兼容 NumPy 的 FFT
• (Prototype) 支持 Nvidia 新一代 A100 GPU 和原生 TF32 格式
• (Prototype) 现在支持 Windows 上的分布式训练
• torchvision
  • (Stable) 变换现在支持张量输入，批量计算，GPU 和TorchScript
  • (Stable) JPEG 和 PNG 格式的原生图像 I/O
  • (Beta) 全新视频读取器 API
• torchaudio
  • (Stable) 增加了对语音记录(wav2letter)、文本到语音(WaveRNN)和源分离(ConvTasNet)的支持

重申一下，从 PyTorch 1.6开始，特性分为 stable（稳定）、 beta（测试） 和 prototype（原型）三个版本。

完整的发布说明可以看这里：

https://github.com/pytorch/pytorch/releases

前端 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] TRANSFORMER 神经网络模块的 C++ 支持

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

性能和性能分析：

[BETA] STACK TRACES 添加到性能分析器

现在，用户不仅可以在性能分析器（profiler）输出表中看到操作符名称/输入，还可以看到操作符在代码中的位置。工作流只需要很少的改变就可以利用这个功能。用户像以前一样使用 autograd profiler，但是带有可选的新参数: with_stack 和 group_by_stack_n。注意: 运行常规性能分析不应该使用这个特性，因为它会增加大量的额外开销。

分布式训练和 RPC：

[BETA] 在 DDP 中支持不均匀的数据集输入

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

[BETA] NCCL 可靠性 - 异步错误/超时处理

在过去，NCCL 的训练运行会因为 collectives 的卡住而无限期地延长，导致用户非常不愉快的体验。如果检测到可能的挂起，这个特性将中止被卡住的 collectives 并抛出异常/中止进程。当与 torchelastic(可以从最后一个检查点恢复训练过程)之类的东西一起使用时，用户可以有更高的可靠性进行分布式训练。这个功能是完全可选的，并且位于需要显式设置才能启用这个功能的环境变量之后(否则用户会看到和以前一样的行为)。

[BETA] TORCHSCRIPT 的RPC_REMOTE 和RPC_SYNC

在之前的版本中，TorchScript 中提供了 torch.distributed.rpc.rpc_async。对于 PyTorch 1.7，此功能将扩展其余两个核心 RPC API: torch.distributed.rpc.rpc_sync 和 torch.distributed.rpc.rpc_sync。这样 TorchScript 中支持的主要 RPC API 就完整了，它允许用户在 TorchScript 中使用现有的 python RPC API (在脚本函数或脚本方法中，它将释放 python GIL) ，并可能提高多线程环境中的应用程序性能。

[BETA] 基于 RPC 的性能分析增强

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

• 在 RPC 上实现了对 TorchScript 函数分析的更好支持
• 实现了与 RPC 一起工作的分析器功能的奇偶校验
• 增加了对服务器端异步 RPC 函数的支持（使用修饰 rpc.functions.async_execution）

用户现在可以使用熟悉的分析工具，比如 torch.autograd.profiler.profile() 和 torch.autograd.profiler.profile()，完全特性支持、配置文件异步函数和 TorchScript 函数在 RPC 框架中是透明的。

[PROTOTYPE] 分布式训练支持 WINDOWS

PyTorch 1.7为 Windows 平台上的 DistributedDataParallel 和 collective 通信带来了原型支持。在这个版本中，支持只包括基于 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])

torchvision：

[STABLE] 变换现在支持张量输入，批量计算，GPU 和 TORCHSCRIPT

torchvision 变换（transforms）现在继承自 nn.Module，可以用 torchscript 编写，应用到 torch 张量输入和 PIL 图像。它们还支持张量 batch 维，并且可以在 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

[STABLE] JPEG 和 PNG 格式的原生图像 IO

torchvision 0.8.0介绍了 JPEG 和 PNG 格式的原生图像读写操作。这些操作符支持 TorchScript 并以 uint8格式返回 CxHxW 张量，因此现在可以成为 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)

[STABLE] RETINANET 检测模型

这个版本为 RetinaNet 添加了一个用ResNet50骨干的预训练模型。

更多的详细信息请查看官网 Blog：

https://pytorch.org/blog/pytorch-1.7-released/