NeurIPS顶会接收，PyTorch官方论文首次曝光完整设计思路

PyTorch 已开源两年多了，现在是最火热的深度学习框架之一。但是，和一些有着论文介绍的开源项目相比，PyTorch 背后的特性和思想还没有被完整介绍过。近日，趁着 NeurlPS 2019 大会即将召开，PyTorch 开发项目组也「水」了一篇论文——PyTorch 框架论文。这篇论文完整且系统地介绍了 PyTorch 本身，引起了社区的关注。

 

目前，这篇论文已经被 NeurlPS 2019 接收为 poster 论文。

 

论文地址：https://papers.nips.cc/paper/9015-pytorch-an-imperative-style-high-performance-deep-learning-library.pdf

 

 

论文放出后，吃瓜群众纷纷流下了感动的泪水，表示：「以后终于有PyTorh的论文可以引用了」，「我们这些年遗漏的参考文献（终于放出来了）。」

 

论文概览

 

近年来，随着深度学习的兴起，机器学习工具也如雨后春笋般发展起来。Caffe、CNTK、TensorFlow、Theano 等很多流行框架都构建了一个表征计算的静态数据流图，这些图可以重复应用于批量数据。

但一直以来，这些深度学习框架要么关注易用性，要么关注速度，很难二者兼顾。Pytorch 的出现打破了这一限制。作为一个机器学习库，它提供了一种命令式、Python 式的编程风格，支持代码作为模型，使得调试变得简单，并且与其他流行的科学计算库保持一致，同时保持高效并支持 GPU 等硬件加速器。

在这篇论文中，作者介绍了驱动 PyTorch 实现的详细原则以及这些原则在 PyTorch 架构中的反映。此外，作者还解释了如何谨慎而务实地实现 PyTorch 运行时的关键组件，使得这些组件能够协调配合，达到令人满意的性能。研究者在几个常见的基准上展示了 PyTorch单个子系统的效率以及整体速度。

 

PyTorch 的诞生背景

 

对于深度学习来说，科学计算领域的四大趋势正变得越来越重要。

 

1. 首先，1960年代，APL、MATLAB、R 和 Julia 等领域特定语言的发展将多维数组（张量）转换为由一组复杂数学算子支持的一级目标，以此对其进行处理。另外，NumPy、Torch、Eigen、Lush 等库的出现使得基于数组的编程在 Python、Lisp、C++、Lua 等通用语言中变得更加高效。

2. 其次，自动微分的发展使得导数的繁冗计算可以完全自动化。autograd 包的出现推动了这一技术在 NumPy 数组中的使用，类似的方法也应用于Chainer、DyNet、Lush、Torch、Jax、Flux.jl 等框架。

3. 再次，随着开源软件运动的发展，科学界已经从 Matlab 等封闭私有软件转向开源 Python 生态，后者包含 NumPy、SciPy、Pandas 等库。

4. 最后，GPU 等通用大规模并行硬件的出现和商业化提供了深度学习方法所需的算力。

 

PyTorch 迎合了这些趋势，它提供了一个由 GPU 加速的、基于数组的编程模型，并通过集成在 Python 生态系统中的自动微分实现可微分。

 

以可用性为中心的设计

 

PyTorch的设计理念相对较新，从易用性、可扩展性的角度进行了设计。

深度学习模型都只是 Python 程序

神经网络从简单的前馈层序列快速演化为非常多样的数值程序，通常由许多循环和递归函数组成。为了适应这一日益增长的复杂性，PyTorch 放弃了基于图-元编程方法的潜在优势，以保持 Python 的命令式编程模型。

PyTorch 将这一特点扩展到了深度学习工作流的所有方面。定义层、构建模型、载入数据、运行优化器、并行训练过程都利用为通用编程开发的熟悉概念来表示。这一解决方案确保任何潜在的新神经网络架构都可以简单地用 PyTorch 实现。


一个简单但完整的神经网络中用作构建块的自定义层。

一个生成对抗网络的简化训练。

 

互操作性和可扩展性

PyTorch 允许与外部库进行双向交换。例如，PyTorch 提供了一种使用 torch.from_numpy() 函数和 .numpy() 张量方法的机制来实现NumPy 数组和 PyTorch 张量使用之间的转换。类似的功能也可用于交换使用 DLPack 格式存储的数据。

此外，许多关键系统都是专门为可扩展性设计的。例如，自动微分系统允许用户为自定义可微分函数添加支持。为此，用户可以定义一个新的 torch.autograd.Function 子类来实现 forward() 和 backward() 方法，这些方法会指定函数及其导数。
 

自动微分

PyTorch 使用算子重载（operator overloading）方法，在每次执行计算函数时构建一个该函数的表征。在其最近的实现中，PyTorch 执行反向模式自动微分，计算有关多元输入的标量输出的梯度。

PyTorch 另一个有趣且不寻常的特性在于，它可以通过在张量上使用突变的代码进行微分，这是命令式程序的基本构建块之一。

 

以性能为中心的实现

 

一个高效的 C++ 核

为了提高性能，PyTorch 的多数代码都是用 C++ 写的。这一核心 libtorch 库用来实现张量数据结构、GPU 和CPU 算子以及基本的并行基元。它还提供了一个自动微分系统，包括用于多数内置函数的梯度公式。

Python 的 bingding 是使用 YAML 元数据文件生成的。这种方法的一个有趣副作用在于，它允许社区快速创建到多个其他语言的 binding ，产生了 NimTorch、hasktorch 等新项目。

 

分离控制和数据流

控制流的解由 Python 和优化的、在主机 CPU 上执行的 C++ 代码来处理，在设备上产生一个算子调用的线性序列。算子可以在 CPU 或 GPU 上运行。

PyTorch 通过利用 CUDA 流机制将 CUDA 内核调用安排到 GPU 硬件 FIFO 来异步执行算子。

 

自定义缓存张量分配器

PyTorch实现了一个自定义的分配器，它递增地构建CUDA内存的缓存并将其重新分配到之后的配额中，而无需进一步使用CUDA API。这种递增的分配对于实现更好的互操作性也非常关键，因为提前占用所有GPU内存会妨碍用户利用其他GPU支持的Python包。为了进一步提高其效率，这一分配器针对深度学习的特定内存使用模式进行了调优。

这种「一流一池（ one-pool-per-stream ）」的设计假设简化了实现，提高了分配器的性能。由于流序列化执行，如果空闲优先于 CPU 上的重新分配，同样的顺序也会发生在 GPU上。因此，只要在与释放的区域相同的流上使用新的分配，分配器就可以立即重新分配在 CPU 上释放的内存。

但这种设计似乎也是有限制的，因为每个流的分配结果是碎片化的，但在实际操作中，PyTorch 几乎从不使用多个流。众所周知，以一种让CUDA 内核协同共享 GPU 的方式来编写 CUDA 内核是非常困难的，因为精确的调度是由硬件控制的。

 

多进程处理

由于全局解释器锁（global interpreter lock，GIL）的 Python 默认实现不允许并行线程进行并行执行，所以为了解决该问题，Python 社区已经建立了一个标准的多进程处理模块，其中包含了大量的实用程序（utility），它们可以使得用户轻易地生成子进程并能够实现基础的进程间通信原语（communication primitive）。

然而，原语的实现使用了与磁盘上持久性（on-disk persistence）相同格式的序列化，这在处理大规模数组时效率不高。所以，PyTorch 将Python 的 multiprocessing 模块扩展为 torch.multiprocessing，这就替代了内置包，并且自动将发送至其他进程的张量数据移动至共享内存中，而不用再通过通信渠道发送。

PyTorch 的这一设计极大地提升了性能，并且弱化了进程隔离（process isolation），从而产生了更类似于普通线程程序的编程模型。

 

引用计数

用户常常设计模型来在训练期间利用所有可用的内存，并且增加批量大小是加速进程的常见方法。所以，为了发挥出色的性能，PyTorch必须将内存视作稀有资源，并小心管理。

在引用计数方面，PyTorch 采取了一种不同的方法：它依赖于一个引用计数方案来追踪每个张量的使用次数，并在该计数为零时立即释放底层内存。需要注意的是，PyTorch 通过集成 Python 自身的引用机制，追踪 libtorch 库内部的引用以及用户在其 Python 代码中所做的外部引用。

需要特别警醒的一点是，我们在已经利用引用计数的语言（CPython、Swift，而非 PyPy 或 Lua 等众多脚本语言）实现，或者在那些允许用户自定义指定、复制和移动行为的语言（如 C++ 和 Rust ）实现中只能保证预期的性能特性。

评估

 

研究者对 PyTorch 和其他几个常用深度学习库的性能进行了比较，发现 PyTorch 在一系列任务上都能实现较突出的性能。所有实验都在一个使用两个英特尔 Xeon E5-2698 v4 CPU 和一个英伟达 Quadro GP100 GPU 的工作站上执行。

 

异步数据流

研究者首先量化了 PyTorch 在 GPU 上异步执行数据流的能力。他们使用内置分析器来度量各种基准，并记录下了单训练步骤上的执行时间线。

下图1展示了 ResNet-50 模型前几步操作执行的典型时间线。在该例中，GPU 执行花费的时间约是 CPU 调度的3倍。精确的比例则取决于主 CPU 和 GPU 的相对性能、每个张量中的组成部件数量以及在 GPU 上实现的浮点运算的平均算法复杂性。

 

图1： Resnet-50模型的前几步操作的轨迹。

 

内存管理

研究者使用英伟达分析器来追踪 CUDA 的执行时间以及 ResNet-50 模型训练迭代期间启动的 CUDA 核心的执行。如下图2所示，首次迭代的行为表现与接下来的迭代截然不同。

 

图2：GPU 上 ResNet-50 模型执行的自动追踪。

 

基准测试

最后，通过与三个流行的图深度学习框架（CNTK、MXNet 和 TensorFlow ）、define-by-run 框架（Chainer）和生产导向型平台（PaddlePaddle）的比较，研究者从整体上得出了 PyTorch 的单机 eager 模式性能。具体结果如下表1所示，PyTorch的性能在最快框架性能的17%以内。

 

表1：AlexNet、VGG-19、ResNet-50、MobileNet、GNMTv2 和 NCF 6 种模型在使用32位浮点运算时的训练速度。

 

PyTorch 的应用

通过计算自 2017年1月PyTorch 首次发布以来各种机器学习工具（包括 Caffe、Chainer、CNTK、Keras、MXNet、Pytorch、TensorFlow 和 Theano）在 arXiv 论文中提及的频率，研究者试图量化机器学习社区对 PyTorch 的接受程度。

如下图3所示，研究者绘制出了所有这些深度学习框架中「PyTorch」每月出现的比例。

 

图3：自 2017 年1 月以来，在所有常见的深度学习框架中，PyTorch 在 arXiv 论文中每月被提及的比例。

 

未来展望

 

除了继续支持深度学习领域最新的趋势和进展之外，研究者计划进一步提升 PyTorch 的速度和可扩展性。最主要的一点是，他们正开发 PyTorch JIT 系列工具，它们可以使得 PyTorch 程序脱离 Python 解释器执行，从而可以得到进一步优化。研究者还打算通过为数据并行化提供高效的原语以及为基于远程过程调用的模型并行化提供 Pythonic 库，进而提升对分布式计算的支持。