跳转到主要内容
博客

通往 1.0 之路:生产就绪的 PyTorch

作者: 2018年5月2日2024年11月16日暂无评论

我们很高兴能为您预览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 vs NHWC vs 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)

脚本模式

追踪模式是最大限度减少对代码影响的好方法,但我们也对那些根本上利用控制流的模型(如RNN)感到非常兴奋。我们的解决方案是脚本模式。

在这种情况下,您编写一个普通的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,并保留其动态特性,这样您就可以使用Caffe2的模型执行器或将模型传输到任何其他支持ONNX的框架,在无Python环境中轻松运行它们。

可用性

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

首先,我们遵循按需付费的原则——如果您不需要优化或导出模型,则无需使用这些新功能,也不会遇到任何缺点。此外,追踪或@script模块/函数的使用可以增量进行。例如,所有这些行为都是允许的:您可以追踪模型的一部分,并在更大的未追踪模型中使用该追踪。您可以将90%的模型用于追踪,并为其中包含一些控制流的子模块使用@script。您可以使用@script编写一个函数,并让它调用一个原生的Python函数。如果@script函数中出现问题,您可以移除注解,代码将在原生Python中执行,在那里可以使用您喜欢的工具和方法轻松调试。可以将追踪和@script视为使用MyPy或TypeScript进行类型注解——每个额外的注解都可以增量测试,并且在您想要优化或生产化之前,都不需要它们。

最重要的是,这些模式将内置于PyTorch的核心中,以便与您现有代码的混合和匹配可以无缝发生。

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

其他更改和改进

生产支持是1.0版本的一大特性,但我们将在标准发布过程中继续优化和修复PyTorch的其他部分。

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

结语

我们的目标是在夏季某个时候发布1.0版本。您可以在拉取请求页面上关注我们的进展。

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