迈向1.0:PyTorch和Caffe2的幸福联姻
编者按:今天是Facebook F8年度开发者大会的第二天,继5月2日发布约会功能、新隐私保护功能、Oculus Go VR头盔等消息后,3日Facebook又带来了机器学习方面的一些重大更新,包括即将在夏季末发布的Pytorch 1.0和新开源的围棋AI项目。以下是论智编译的Pytorch 1.0公告。
Pytorch源:pytorch.org/2018/05/02/road-to-1.0.html
Caffe2源:caffe2.ai/blog/2018/05/02/Caffe2PyTorch1_0.html
Caffe2公告摘要
构建和包
在Pytorch 1.0中,Caffe2和PyTorch的python软件包(在pip和conda中)会合并;
继续提供本地库和python扩展作为单独的安装选项(和Caffe2和PyTorch现状一样);
所有交叉编译的构建模式均支持Caffe2平台(iOS,Android,Raspbian,Tegra等),将继续拓展新平台支持。
模型生成(前端)
Caffe2的图形构建API(brew,core.Net)继续保留,现有的序列化模型NetDefs会拥有向后兼容性;
nn.Module-like将成为构建神经网络的首选。Pytorch 1.0会用“混合前端”取代现有前端,它利用跟踪和编译功能,能以图形格式(兼容Caffe2的NetDef和ONNX)提取完全序列化的模型,可用于高效部署。
Caffe2和PyTorch的运算符将合并使用;
PyTorch 1.0能本地支持ONNX,连接其他框架或库以便提取、传递模型。
精简和整合
Caffe2高度可扩展的执行引擎(各种加速后端和库集成和可伸缩图执行器)基本保持不变,还能与Python程序段交互操作,以实现快速原型设计;
Caffe2现有的predictor支持会成为本地模型在加速硬件支持上的主要手段(数据中心和移动设备可用)。
硬件集成和加速库支持
在PyTorch 1.0软件包的原型环境中直接提供Caffe2的各种设备支持和运行集成;
保留Caffe2针对加速硬件或完整图形运行的大部分接口代码,现有集成仍然有效。对于图形运行,官方正致力于ONNX的可视化。
Pytorch公告全文
亲爱的Pytorch用户,
在这篇文章中,我们想带你预览一下Pytorch 1.0的发展路线。Pytorch 1.0是Pytorch的下一代版本,在过去一年中,我们陆续推出了PyTorch0.2、0.3和0.4,并把原先类似[Torch + Chainer]的接口转换成了更简洁的形式,我们还引入double-backward、类似numpy的函数、高级索引,同时删除了变量样板。截至目前,我们确信所有API都处于稳定合理的状态,因此现在也是时候可以放心发布Pytorch 1.0了。
但是,Pytorch 1.0不仅仅意味着更高的稳定性。
Pytorch最大的优势之一是一流的Python集成水平,强制性的风格、简单的API选项,这些都是它方便用户研究和操作的重要因素。但它也有不少缺点,其中最大的一点是“production-support”。这意味着为了更高效地大规模运行神经网络,用户必须对模型进行大量处理:
做大型项目时,用户需要把模型导出到C++;
要针对iPhone、Android、Qualcomm等其他移动设备系统和硬件进行优化;
为了更快地推断,用户需要使用更高效的数据布局和执行内核融合(节省10%地耗时和内存占用已经称得上是个巨大胜利);
量化精度(如8-bit推断)。
许多初创公司、大公司以及一些想用PyTorch构建产品的个人用户都曾向我们要求提供产品支持。在Facebook(PyTorch最大的利益相关者),我们已经有了Caffe2这个面向移动设备的平台,它曾在超过10亿部手机上运行,横跨8代iPhone和6代Android CPU架构。为了能更好地支持Intel / ARM、TensorRT,它也优化过服务器推断。考虑到PyTorch团队想开发的内容和Caffe2已经成熟的功能基本一致,因此我们决定结合PyTorch和Caffe2,让后者为PyTorch提供生产支持。
有了它,研究人员和终端用户就能在不产生可用性问题的基础上进行产品设计,而这需要创造性的解决方案。
生产!= 研究人员的痛点
提高PyTorch的生产能力意味着要引入更复杂的API和更多的可配置选项的数量。详细来说,就是一个配置内存布局(NCHW vs NHWC vs N、C/32、H、W、32,每个提供不同的性能特性),量化精度(8-bit?3-bit?),低水平内核融合(把Conv + BatchNorm + ReLU融合进一个内核),以及分离后端选项(把MKLDNN作为某几层的后端,NNPACK作为剩余层的后端)等。
PyTorch的核心目标是为研究和操作提供一个良好的平台,因此除了以上这些优化,我们也一直致力于设计一个严格的约束,以避免因完善功能牺牲用户体验。
为了做到这一点,我们这次推出了torch.jit
——一个即时(JIT)编译器,它能在运行时重写用户模型并实现快速产品化。这个新的JIT编译器还能把模型迁移到C++上,然后以Caffe2的比特精度运行。
对于Pytorch 1.0,你的代码可以持续运行,但我们没有对已有API做过多调整。
当然,这是一个可选项。用户可以使用torch.jit编译器把自己的模型导出到无Python环境并提高其性能。下面让我们对它做一下详细介绍。
torch.jit:一个为模型设计的即时编译器
我们坚信一点,如果直接用Python代码写一个模型来应用,它的适用性太低了。这也是PyTorch为什么这么灵活的原因,但有得必有失,如果用户不下指令,PyTorch就不知道接下来该运行哪一段。这对生产和自动性能优化来说是个不小的问题,因为在执行代码前,它们需要获知模型的整体情况。
为了解决这个问题,我们提供了两种从代码中恢复信息的方法,一种基于追踪本地Python代码,另一种则基于编译表示中间形态的Python语言子集。经过充分讨论后,我们认为它们在不同环境中都由突出表现,因此用户可以搭配使用。
追踪模式
PyTorch tracer是一个函数torch.jit.trace
,它记录了当前代码区执行的所有本地PyTorch操作和它们之间的数据依赖关系。其实早在PyTorch 0.3时我们就提供了类似功能,它能被用于通过ONNX导出模型。而Pytorch 1.0的变化在于你可以在C++运行过程中就直接执行追踪,而不需要到别处。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)
script模式
追踪是一种好方法,但我们希望能为一些涉及大量循环的模型,如RNN,提供同样优秀的解决方案。这就是我们接着要介绍的script模式。
在script模式下,除了无法使用一些特殊的语言功能,用户只需写一个常规的Python函数。一旦你隔离其他函数,用@script
标记出你想要编译的函数,Pytorch 1.0就会为这个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模式,Pytorch 1.0输出的都会是一个非Python的模型表示。它可用于模型优化,或从python中导出模型用于生产环境。
像这种将大部分模型信息提取为中间表示的做法有助于执行复杂的整体程序优化,在这个背景下,我们可以把计算分配给专门用来计算图形的AI加速器,以提高整体效率。事实上我们也正在这么做,我们已经在尝试融合GPU操作提高小型RNN的性能,成效斐然。
它还意味着用户能使用Caffe2现有的高性能后端高效运行模型。此外,用@script标记的函数和模块还可以以动态的方式导出到ONNX,这就意味着它们可以轻松地在非Python语言环境中运行。
可用性
可用性是我们十分关注的一个点,我们知道,如果不是直接在Python里运行代码,模型后期调试起来会很麻烦,但我们希望能做更多的事,而不是把用户局限在一种编程语言上。
首先,这是一项付费服务——如果用户不需要优化和导出模型,他们就不用在意这些新功能,也不会看到其中的缺陷。其次,追踪模式和script模式是可以逐步使用的,例如用户可以追踪部分模型、在大型未追踪模型中嵌入追踪模式、对模型的90%使用追踪、只对循环代码使用@script……他们可以用@script编写一个函数,并让它调用一个本地python函数,如果@script函数中出现错误,他们也可以删除注释。总而言之,用户喜欢怎么调试,就可以怎么调试。
最重要的是,这些模式将被构建到PyTorch的核心中,以便与用户现有代码混合并无缝运行。
其他更新与改进
“production-support”是Pytorch 1.0最重要的特征,但作为标准发布过程的一部分,我们也会继续优化和修复PyTorch的其他内容。
在后端,PyTorch的一些更新可能会影响用户编写C和C ++的扩展。因为我们正在替换(或重构)后端ATen库,以整合来自Caffe2的功能和优化。
进度观察地址:github.com/pytorch/pytorch/pulls