由 PyTorch 团队编写

我们想为您展示 PyTorch 1.0 的路线图预览,这是 PyTorch 的下一个版本。在过去的一年里,我们发布了 0.2、0.3 和 0.4 版本,将 PyTorch 从类似[Torch+Chainer]的接口转变为更简洁的界面,增加了双向反向传播、类似 NumPy 的函数、高级索引,并移除了 Variable 模板代码。目前,我们相信 API 已经处于合理且稳定的状态,可以放心发布 1.0 版本。

然而,1.0 版本并不仅仅关乎接口的稳定性。

PyTorch 最大的优势之一是其一流的 Python 集成、命令式风格、API 的简洁性和选项。这些特点使得 PyTorch 非常适合研究和易于开发。

它最大的缺点之一是生产支持。我们所说的生产支持是指为了使模型在大规模上高效运行,必须对模型进行无数次的调整:

  • 将代码导出到仅限 C++运行时,用于大型项目
  • 优化 iPhone、Android、高通和其他系统上的移动系统
  • 使用更高效的数据布局和执行内核融合以实现更快的推理(在规模上节省 10%的速度或内存是一个巨大的胜利)
  • 量化推理(例如 8 位推理)

创业公司、大型企业和任何想要围绕 PyTorch 构建产品的人士都要求生产支持。在 Facebook(PyTorch 的最大股东)我们拥有 Caffe2,它一直是我们的生产就绪平台,在我们的数据中心运行,并已发货至超过 10 亿部手机,覆盖了八代 iPhone 和六代 Android CPU 架构。它支持在 Intel / ARM 上的服务器优化推理、TensorRT 支持以及所有必要的生产组件。考虑到 PyTorch 团队与该平台紧密合作所积累的所有价值,我们决定将 PyTorch 与 Caffe2 结合,从而为 PyTorch 提供生产级就绪性。

在不影响我们研究人员和最终用户易用性的情况下支持生产功能需要创造性的解决方案。

生产不是研究人员的痛苦。

添加生产能力涉及增加 API 的复杂性和模型的配置选项数量。用户可以配置内存布局(NCHW、NHWC、N,C/32,H,W,32,每种都提供不同的性能特性)、量化(8 位?3 位?)、融合低级内核(您使用了 Conv + BatchNorm + ReLU,让我们将它们融合成一个内核)、独立的后端选项(MKLDNN 后端用于几层,NNPACK 后端用于其他层)等。

PyTorch 的核心目标是提供一个出色的研究平台和可玩性。因此,当我们添加所有这些优化时,我们一直在努力确保不会以可用性为代价。

为了实现这一点,我们引入了 torch.jit ,这是一个即时编译器(JIT),在运行时将您的 PyTorch 模型重写以实现生产效率。JIT 编译器还可以将您的模型导出,以便在仅基于 Caffe2 位的 C++运行时中运行。

在 1.0 版本中,您的代码将继续按原样工作,我们不会对现有的 API 进行任何重大更改。

将您的模型制作成生产就绪状态是一个可选的注释,它使用 torch.jit 编译器将您的模型导出到无 Python 环境,并提高其性能。让我们详细了解一下 JIT 编译器。

torch.jit :为您的模型提供的 JIT 编译器

我们坚信,直接将模型指定为惯用的 Python 代码所获得的效率很难匹敌。这正是 PyTorch 如此灵活的原因,但也意味着 PyTorch 几乎永远不会知道您接下来要运行的操作。然而,这却是导出/生产化的一个重大障碍,因为它们需要在执行之前就完全了解计算过程。

我们提供了两种可选的方法来从您的代码中恢复这些信息,一种基于跟踪原生 Python 代码,另一种基于编译 Python 语言中注释的子集到一个无 Python 的中间表示形式。经过充分讨论,我们得出结论,它们在不同的环境中都将非常有用,因此您将能够自由地混合和匹配它们。

追踪模式

PyTorch 的追踪器, torch.jit.trace ,是一个记录代码区域中所有原生 PyTorch 操作及其之间数据依赖关系的函数。实际上,PyTorch 从 0.3 版本开始就拥有追踪器,它被用于通过 ONNX 导出模型。现在发生变化的是,你不再需要将追踪结果发送到其他地方执行——PyTorch 可以为你重新执行它,使用精心设计的、高性能的 C++运行时。随着 PyTorch 1.0 的开发,这个运行时会集成 Caffe2 提供的所有优化和硬件集成。

这种方法的最大的好处是它并不真正关心你的 Python 代码是如何结构的——你可以追踪生成器或协程、模块或纯函数。由于我们只记录原生 PyTorch 操作符,这些细节对记录的追踪没有影响。然而,这种行为是一把双刃剑。例如,如果你的模型中有一个循环,它将在追踪中展开,循环体将被复制运行次数,就像循环运行了多少次一样。这为无成本的抽象(例如,你可以遍历模块,而实际的追踪将没有循环开销!)提供了机会,但另一方面,这也会影响数据相关的循环(例如,处理长度不同的序列),实际上将单个长度硬编码到追踪中。

对于不包含循环和 if 语句的网络,追踪是非侵入性的,并且足够健壮,可以处理各种编码风格。以下代码示例展示了追踪的样子:

# This will run your nn.Module or regular Python function with the example
# input that you provided. The returned callable can be used to re-execute
# all operations that happened during the example run, but it will no longer
# use the Python interpreter.
from torch.jit import trace
traced_model = trace(model, example_input=input)
traced_fn = trace(fn, example_input=input)

# The training loop doesn't change. Traced model behaves exactly like an
# nn.Module, except that you can't edit what it does or change its attributes.
# Think of it as a "frozen module".
for input, target in data_loader:
    loss = loss_fn(traced_model(input), target)

脚本模式

追踪模式是一种减少对代码影响的好方法,但我们也非常兴奋于那些根本利用控制流,如 RNNs 的模型。我们针对此问题的解决方案是脚本模式。

在这种情况下,你将编写一个常规的 Python 函数,但不能再使用某些更复杂的高级语言特性。一旦你隔离了所需的功能,你可以通过使用 @script 装饰器来让我们知道你希望该函数被编译。这个注解将直接将你的 Python 函数转换为我们的高性能 C++运行时。这使得我们可以恢复所有 PyTorch 操作,包括循环和条件语句。它们将被嵌入到我们对该函数的内部表示中,并在每次运行该函数时被考虑在内。

from torch.jit import script

@script
def rnn_loop(x):
    hidden = None
    for x_t in x.split(1):
        x, hidden = model(x, hidden)
    return x

优化和导出

无论你使用追踪还是 @script ,结果都是模型的无 Python 表示,可以用来优化模型或从 Python 导出模型以用于生产环境。

将模型的大段提取到中间表示形式,使得进行复杂的整体程序优化和将计算卸载到专门的操作图 AI 加速器成为可能。我们已经在开发这些优化的初步工作,包括将 GPU 操作融合在一起以提升较小 RNN 模型的性能。

这也允许我们利用 Caffe2 今天已有的高性能后端来高效运行模型。此外,@script 函数(以及模块!)可以完全导出到 ONNX 中,保留其动态特性,这样您就可以在无需 Python 的环境中轻松运行它们,使用 Caffe2 的模型执行器或通过将模型转移到任何支持 ONNX 的其他框架中。

可用性

我们非常重视保持当前的可用性水平,我们知道不在 Python 中直接执行代码会导致调试更加困难,但我们一直在思考这个问题,并确保您不会被锁定在完全不同的编程语言中。

首先,我们遵循按使用付费的原则——如果您不需要优化或导出模型,您无需使用这些新功能,也不会看到任何缺点。此外,使用可追踪或@script 模块/函数可以逐步进行。例如,以下行为都是允许的:您可以追踪模型的一部分,并在更大的非追踪模型中使用追踪。您可以使用追踪覆盖模型中的 90%,而使用@script 处理实际具有控制流的子模块。您可以使用@script 编写函数,并让它调用原生 Python 函数。如果在@script 函数中出现问题,您可以移除注释,代码将在原生 Python 中执行,使用您喜欢的工具和方法进行调试变得容易。将追踪和@script 视为使用 MyPy 或 TypeScript 进行类型注解——每个额外的注解都可以逐步测试,直到您想要优化或生产化之前,无需任何注解。

最重要的是,这些模式将集成到 PyTorch 的核心中,以便与现有代码无缝混合和匹配。

注意:这些组件的 JIT 名称有些名不副实,这源于历史原因。PyTorch 中的跟踪/函数执行最初是一个优化 JIT 编译器,它生成融合的 CUDA 内核,但后来扩展到包括优化、@script 和导出。当它准备发布时,我们可能会将此功能重命名为混合前端,但我们希望在这里按照代码中的名称展示它,以便您在开发过程中可以跟随。

其他更改和改进

生产支持是 1.0 版本的主要功能,但我们将继续优化和修复 PyTorch 的其他部分,作为标准发布流程的一部分。

在后端方面,PyTorch 将看到一些变化,这可能会影响用户编写的 C 和 C++扩展。我们正在替换(或重构)后端 ATen 库,以纳入 Caffe2 的功能和优化。

最后一句话

我们计划在夏季某个时候发布 1.0 版本。您可以在“拉取请求”页面跟踪我们的进度。

您可以从 Caffe2 项目的角度阅读此内容:https://caffe2.ai/blog/2018/05/02/Caffe2_PyTorch_1_0.html