整合PyTorch 0.4和Caffe 2，PyTorch 1.0能挑战TensorFlow吗？

Posted 2021-05-02 AI科技大本营



译者 | 梁红丽

编辑 | Mavis

【AI 科技大本营导读】5月2日，在加利福尼亚州举办的年度开发者 F8 大会上，Facebook 正式推出 PyTorch 1.0 。其实，早在 2017 年 1 月，Facebook 就首次公布了该信息，截至目前，它已被下载超过 110 万次，是过去一个月研究门户网站 Arxiv 上的第二大深度学习框架，排名第一的是 TensorFlow 。

另外，手握 ArXiv Sanity 大数据的特斯拉人工智能部门主管 Andrej Karpathy 给出的精确数据显示，过去一个月，各个框架在论文中被提到（单次计算）的比例分别是：TensorFlow 14.3% ，PyTorch 4.7％ ，Keras 4.0％ ，Caffe 3.8％ ，Theano 2.3％ ，Torch 1.5％ ，MXNet 、Chainer 和 CNTK 均小于 1％。

2017 年 3 月，Karpathy 第一次做这个全面统计的时候，各框架的排名是：



对比两组数据可以发现，PyTorch 涨势惊人，看来想要挑战 TensorFlow ，并不是没有可能。接下来，和营长一起了解下 PyTorch 1.0 。

过去几年中，Facebook 发行了 0.2 , 0.3 , 0.4 几个版本，从类似 Torch+Chainer 的界面改组为试验版本，增加了双击后退、numpy-like 函数、高级索引和删除变量功能，且界面更清晰。

但是，1.0 的界面不仅仅是稳定性有所提高。

PyTorch 最大的优势之一就是一流的 Python 交互、命令式风格、API 和选项简约，所有这些特点都使 PyTorch 利于研究和整改。

PyTorch 最大的缺点一直以来都是产品支持。“产品支持”指为了使模型在大规模使用时高效运行而必须对其做出的不计其数的修改，包括：

• 在大项目中，将模型输出到只支持 C++ 的环境中使用；

• 优化 iPhone，android，Qualcomm 和其他系统的移动端系统；

• 用更高效的数据框架和内核加快推理（大规模使用中速度加快 10% 或内存减少 10% 就是重大胜利）。

• 量化推理

初创公司、大公司和任何想用 PyTorch 做产品的人都要求提供产品支持。在 Facebook，我们有 Caffe2，它随时可以投产，在数据中心运行、被载入超过 10 亿部手机中，包括八代的 iPhone 和六代的 Android CPU 架构。它有 Intel/ARM 的服务器优化推理，TensorRT 支持和所有产品生产所必须的条件。考虑到所有这些有价值的特点都集成在与 PyTorch 团队工作如此贴近的平台上，我们决定将 PyTorch 和 Caffe2 进行嫁接，使 PyTorch 可直接用于产品。

▌产品 != 研究者的痛

增加产品功能需要提高 API 的复杂度和模型设置选项的数量。模型设置项有内存配置（NCHW vs NHWC vs N,C/32, H, W,32，每个的表现都不同）、量化 bit 数（8 bit 或 3 bit？）、用于组合的低层核（如将 Conv+BatchNorm+ReLU 组合为单核）、可分层选择的后端（例如几层用MKLDNN，另外几层用 NNPACK）等。

PyTorch 的中心目标是为研究和可修改性提供良好平台，因此，加入这些优化的同时，我们也进行了硬件设计，以使这些优化不影响使用。

为了完成这一目标，我们引入了 torch.jit ——一个即时（ JIT ）编译器，它在运行时使用 PyTorch 模型并将其重写，使其以产品效率运行。即时编译器还可以将模型输出到基于 Caffe2 bit 的 C++ 环境。

1.0 版本中，你的代码可一如既往地使用，我们对当前 API 没有做大的改动。

▌Torch.jit：模型的即时编译器

我们深知将具体模型直接写为 Python 代码来满足效率并不容易，这是 PyTorch 如此灵活的原因，但这也意味着 PyTorch 很可能不知道你下一步要运行什么。这会阻碍输出/产品化和重量级自动性能优化，因为它们在执行前就需要如何进行计算的先验知识。

我们提供了从代码还原该信息的两种方法，一种以追踪本地 Python 代码为基础，一种以编译  Python 语言子集并注释为没有 Python 的中间表示为基础。深入讨论后，我们得出结论：它们均可在不同环境中使用，所以你可以自由组合。

 

▌追踪模式

PyTorch 追踪器是一个记录在代码中执行的所有 PyTorch 本地操作及操作间数据依赖的函数。和之前的版本不同，在 1.0 版本中，不需要再记录轨迹放在其他地方运行，PyTorch 会代你用认真设计的表现性能好的 C++ 环境重新执行。

这种方法的最大好处是，它不关心你的 Python 代码是如何组织的，你可以通过生成器或协成程序（ coroutines ）、模块或函数组进行追踪。

对不包含循环和判断的网络，追踪是非侵入式的，且处理多种多样的代码风格时足够稳定。下面的代码提供了追踪示例：

# 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++ 环境。

 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 模块/函数可以递进使用。

最重要的是，这些模式将建于 PyTorch 的核心部分，这样可将它们与现有代码进行无缝对接和匹配。

▌结语

产品支持是 1.0 的特点，但我们会在标准发行的同时继续优化和改进 PyTorch 的其他部分。

在后端部分，PyTorch 会略有改动，这可能影响用户编写的 C 或 C++ 的拓展。我们正在替换后端 ATen library，旨在兼容 Caffe2 的特点和优化操作。

原文链接：

https://pytorch.org/2018/05/02/road-to-1.0.html

招聘

AI科技大本营现招聘AI记者和资深编译，有意者请将简历投至：[e-mail]，期待你的加入！

AI科技大本营

公众号ID：rgznai100

☟☟☟点击 | 阅读原文 | 查看更多精彩内容