Paper | PyTorch: An Imperative Style, High-Performance Deep Learning Library

目录

• 0. 摘要
• 1. 简介
• 2. 背景
• 3. 设计原则
• 4. 针对易用性的核心设计
  • 4.1 让深度学习模块不过是Python程序
  • 4.2 互用性和可拓展性
  • 4.3 自动差分
• 5. 针对高性能的PyTorch实现
  • 5.1 高效的C++核
  • 5.2 分而治之
  • 5.3 定制的张量分配缓存
  • 5.4 多进程
  • 5.5 Reference counting

终于可以引用PyTorch了！

0. 摘要

深度学习框架要么关注易用性，要么关注速度，很少同时关注二者。但PyTorch作为一个机器学习库，同时做到了这两点，说明了这两点是兼容的。

PyTorch提供了 基础必要 且 Pythonic 的编程风格，让调试更简单，并且支持许多科学计算库，同时提供硬件加速例如GPUs。

作者向我们展示：他们如何让PyTorch的关键组分协同工作的，并且提供了PyTorch实现的原则细节。作者强调：PyTorch就是一个常规的Python程序，并且完全由用户控制。

本文还在一些benchmarks测试了PyTorch的速度，并且展示了PyTorch每个独立子系统的效率。

1. 简介

许多框架，例如Caffe，TensorFlow和Theano，都是构建静态数据流图（static dataflow graph）来表示计算流程。

该静态图可以被数据的不同batch重复使用。静态图很好，可以提前让我们预知计算流程，并且可以在理论上用于提升性能和可拓展性。但另一方面，静态图的使用降低了易用性、调参体验，限制了能够被表示的计算类型。

因此，另一些框架（如Chainer）尝试用动态图机制（dynamic eager execution）处理深度学习任务。缺点是要么牺牲性能，要么使用更生涩的语言（如Torch等）。

然而，只要我们处理得当，动态图机制能够克服这些缺点。本文即介绍PyTorch。PyTorch是一个Python库，可以快速执行动态张量运算，并且配备了自动差分和GPU加速技术。最重要的是，其性能可比肩当前最快的库。

2. 背景

在科学计算领域，有4个面向深度学习的主要发展趋势：

1. 基于数组的编程更高效。从1960年代开始，各领域专用的语言，例如MATLAB和R，将张量转换为一等对象（first-class object），再对这些张量进行操作。而这些一等对象由全面的数学算子提供支持。

2. 自动差分技术。这项技术让用户从令人生畏的梯度计算中解脱出来。主要归功于autograd包。

3. 开源的Python生态系统。这里有NumPy，SciPy，Pandas等。也正因如此，2014伊始，越来越多的框架支持Python界面。

4. 硬件加速。普适的并行硬件越来越普及，例如GPUs。还有cuDNN等专用库提供支持，使得越来越多的平台 实现并受益于 硬件加速。

作者用一句话总结了以上趋势，从而解释了PyTorch的优越性：

PyTorch builds on these trends by providing an array-based programming model accelerated by GPUs and differentiable via automatic differentiation integrated in the Python ecosystem.

3. 设计原则

1. Pythonic。PyTorch应该成为Python生态系统的一等对象。PyTorch应与Python保持连贯性和一致性，整合标准的画图、调试和数据处理工具。

2. 用户至上。PyTorch应该把复杂的机器学习内核包裹起来，并提供没有性能损失的外接APIs。

3. 保证易用性和兼容性。可以牺牲10%的速度，换取更好的易用性。此外，用户也可以不用预先提供的库，而自己编写代码以达到更高的性能。

4. 不追求实现上的完美，而追求实用性。对于PyTorch的内部实现，我们不一定要追求完美；相反，为了适应AI的快速变化，为了能够搭载未来的额外的特征，我们可以简化内部实现，而牺牲一定的性能。

4. 针对易用性的核心设计

4.1 让深度学习模块不过是Python程序

机器学习发展迅速，从手写数字识别到星际争霸游戏。因此，神经网络也在不断发展，从简单的前向网络堆叠，到各式各样的复杂结构（例如包含循环和回路）。

为了适应这种快速变化，PyTorch放弃了基于图-元编程的方法（graph-metaprogramming based approach），而是保持了Python的必要编程结构。这种思路最早引用于Chainer和Dynet，而PyTorch将其拓展至深度学习的各个工作环节，例如定义网络、模型组建、数据装载、优化器运行和并行化训练。

例如我们最熟悉的layer。PyTorch将其表达为Python的类。其constructors创建并初始化其参数，forward方法处理输入激励。

但要强调的是，PyTorch并不强迫用户使用上述代码。

由于PyTorch程序是动态加载的，因此所有特征在设计过程中都是可触及的。例如打印状态，标准调试器和通用的可视化工具（如matplotlib）。用户在运行程序前，不再需要等待冗长的编译过程。最重要的是，用户可以实时观察程序的执行过程。

4.2 互用性和可拓展性

这一点使得PyTorch可以拥抱Python的生态系统。例如，PyTorch张量和NumPy数组间可以互相转换。注意，这种互换过程不涉及数据复制；二者仅仅表征对同一共享的内存区域的各自解释（objects on both sides only describe how to interpret a memory region which is shared among them）。因此，无论转换的数组有多大，花费的时间是恒定的。

不仅如此，许多核心子系统都被设计为可拓展的。例如自动差分系统，可以支持用户自定义的函数。为了实现这一点，用户只需要定义torch.autograd.Function的一个新子类，并且包含forward()和backward()方法。类似的，用户还可以自定义新的数据集，而只需要继承torch.utils.data.Dataset并且包含__getitem__和__len__方法。而具体实现方法完全由用户决定，可以随意DIY。

用户如果觉得PyTorch的哪一个模块无法达到要求，可随意更换之。

4.3 自动差分

首先，自动差分是非常必要的，因为目前的深度学习方法大多依赖于基于梯度的优化方法。然而，Python是一种动态编程语言，在运行时允许大多数行为的变化。这与预先的差分是冲突的。

为了解决这一问题，PyTorch采用了操作过载方法（operator overloading approach），即在每一次执行运算时，建立关于运算的表示。

最后，哪怕在运行过程中张量有突变，差分也能够完成。为了保证安全性，我们为张量设置了版本控制系统（versioning system）。在突变发生时，为了避免不易察觉的性能损失，我们通常会告知用户错误发生，而不是自动予以纠正。

5. 针对高性能的PyTorch实现

用Python解释器跑深度学习算法是出了名的麻烦：例如，全局锁（global interpreter lock）会保证任意时刻只有一个进程在运行（而避免多线程）。对于那些采用静态图的框架，它们通过定制的解释器来避免这个问题。

而PyTorch从另一个角度解决了这一问题。

5.1 高效的C++核

PyTorch的大部分都是用C++实现的，以保证高性能。例如libtorch库用来表示张量的数据结构、CPU和GPU操作以及基本的并行算子。还有自动差分系统等。这样，我们就可以实现多线程执行PyTorch的核心操作。

这带来了一个额外的好处。当我们认为Python不行时，我们还可以用C++的一等bindings。

5.2 分而治之

PyTorch将control（如程序回路、分支）和数据流（如张量及其操作）严格分开。控制流通过Python和C++代码控制，并在CPU上执行；算子既可以在CPU上运行，也可以在GPU上运行。

5.3 定制的张量分配缓存

几乎每一个算子都要动态分配一个输出张量。因此，针对这种动态内存分配的速度优化就显得尤为重要。PyTorch对此进行了优化。

PyTorch还对内存分配进行了优化，例如要求每次分配512字节倍数的内存。

还有很多，参见5.3节。

5.4 多进程

为了解决Python的全局锁问题，Python社区建立了标准的多进程模块。PyTorch进一步将其发展为torch.multiprocessing模块。

5.5 Reference counting

为了最大化效率，用户通常会用上所有GPU内存。因此PyTorch的内存管理显得尤为重要。

对于某张量，动态语义库可能无法事先知晓其未来会被如何使用。因此，垃圾回收就成了一种常用的内存管理策略。PyTorch采用了一种特别的策略，快速释放空闲GPU内存。详情参见5.5节。

实验略。