Edward Z. Yang作者panda参与

万字综述,核心开发者全面解读PyTorch内部机制

斯坦福大学博士生与 Facebook 人工智能研究所研究工程师 Edward Z. Yang 是 PyTorch 开源项目的核心开发者之一。他在 5 月 14 日的 PyTorch 纽约聚会上做了一个有关 PyTorch 内部机制的演讲,本文是该演讲的长文章版本。

大家好!今天我想谈谈 PyTorch 的内部机制。

这份演讲是为用过 PyTorch并且有心为 PyTorch 做贡献但却被 PyTorch 那庞大的 C++ 代码库劝退的人提供的。没必要说谎:PyTorch 代码库有时候确实让人难以招架。

本演讲的目的是为你提供一份导航图:为你讲解一个「支持自动微分的张量库」的基本概念结构,并为你提供一些能帮你在代码库中寻路的工具和技巧。我预设你之前已经写过一些 PyTorch,但却可能还没有深入理解机器学习软件库的编写方式。

本演讲分为两部分:在第一部分中,我首先会全面介绍张量库的各种概念。我首先会谈谈你们知道且喜爱的张量数据类型,并详细讨论这种数据类型究竟能提供什么,这能让我们更好地理解其内部真正的实现方式。

如果你是一位 PyTorch 高级用户,你可能已经熟悉其中大部分材料了。我们也会谈到「扩展点(extension points)」的三个概念、布局(layout)、设备(device)和数据类型(dtype),这能引导我们思考张量类的扩展的方式。在 PyTorch 纽约聚会的现场演讲中,我略过了有关自动梯度(autograd)的幻灯片,但我在这里会进行一些讲解。

第二部分会阐述真正用 PyTorch 写代码时所涉及的基本细节。我会告诉你如何在 autograd 代码中披荆斩棘、什么代码是真正重要的以及怎样造福他人,我还会介绍 PyTorch 为你写核(kernel)所提供的所有炫酷工具。

概念

张量

张量是 PyTorch 中的核心数据结构。对于张量直观上所表示的东西,你可能已有很好的理解:张量是一种包含某种标量类型(比如浮点数和整型数等)的 n 维数据结构。我们可以将张量看作是由一些数据构成的,还有一些元数据描述了张量的大小、所包含的元素的类型(dtype)、张量所在的设备(CPU 内存?CUDA 内存?)

另外还有一个你可能没那么熟悉的元数据:步幅(stride)。stride 实际上是 PyTorch 最别致的特征之一,所以值得稍微多讨论它一些。

张量一个数学概念。但要在我们的计算机中表示它,我们必须为它们定义某种物理表示方法。最常用的表示方法是在内存中相邻地放置张量的每个元素(这也是术语「contiguous(邻接)」的来源),即将每一行写出到内存,如上所示。在上面的案例中,我已经指定该张量包含 32 位的整型数,这样你可以看到每一个整型数都位于一个物理地址中,每个地址与相邻地址相距 4 字节。为了记住张量的实际维度,我们必须将规模大小记为额外的元数据。

所以这幅图与步幅有什么关系?

假设我想要读取我的逻辑表示中位置张量 [0,1] 的元素。我该如何将这个逻辑位置转译为物理内存中的位置?步幅能让我们做到这一点:要找到一个张量中任意元素的位置,我将每个索引与该维度下各自的步幅相乘,然后将它们全部加到一起。在上图中,我用蓝色表示第一个维度,用红色表示第二个维度,以便你了解该步幅计算中的索引和步幅。进行这个求和后,我得到了 2(零索引的);实际上,数字 3 正是位于这个邻接数组的起点以下 2 个位置。

(后面我还会谈到 TensorAccessor,这是一个处理索引计算的便利类(convenience class)。当你使用 TensorAccessor 时,不会再操作原始指针,这些计算过程已经为你隐藏了起来。)

步幅是我们为 PyTorch 用户讲解方法的基本基础。举个例子,假设我想取出一个表示以上张量的第二行的张量

使用高级的索引支持,我只需写出张量 [1, :] 就能得到这一行。重要的是:当我这样做时,不会创建一个新张量;而是会返回一个基于底层数据的不同域段(view)的张量。这意味着,如果我编辑该视角下的这些数据,它就会反映在原始的张量中。

在这种情况下,了解如何做到这一点并不算太困难:3 和 4 位于邻接的内存中,我们只需要记录一个说明该(逻辑)张量的数据位于顶部以下 2 个位置的偏移量(offset)。(每个张量都记录一个偏移量,但大多数时候它为零,出现这种情况时我会在我的图表中省略它。)

演讲时的提问:如果我取张量的一个域段,我该如何释放底层张量的内存?

答案:你必须制作该域段的一个副本,由此断开其与原始物理内存的连接。你能做的其它事情实际上并不多。另外,如果你很久之前写过 Java,取一个字符串的子字符串也有类似的问题,因为默认不会制作副本,所以子字符串会保留(可能非常大的字符串)。很显然,Java 7u6 将其固定了下来。

如果我想取第一列,还会更有意思:

当我们查看物理内存时,可以看到该列的元素不是相邻的:两者之间有一个元素的间隙。步幅在这里就大显神威了:我们不再将一个元素与下一个元素之间的步幅指定为 1,而是将其设定为 2,即跳两步。(顺便一提,这就是其被称为「步幅(stride)」的原因:如果我们将索引看作是在布局上行走,步幅就指定了我们每次迈步时向前多少位置。)

步幅表示实际上可以让你表示所有类型的张量域段;如果你想了解各种不同的可能做法,请参阅 https://ezyang.github.io/stride-visualizer/index.html

我们现在退一步看看,想想我们究竟如何实现这种功能(毕竟这是一个关于内部机制的演讲)。如果我们可以得到张量的域段,这就意味着我们必须解耦张量的概念(你所知道且喜爱的面向用户的概念)以及存储张量的数据的实际物理数据的概念(称为「存储(storage)」):

也许会有多个张量共享同一存储。存储会定义张量的 dtype 和物理大小,同时每个张量还会记录大小、步幅和偏移量,这定义的是物理内存的逻辑解释。

有一点需要注意:总是会存在一个张量-存储对,即使并不真正需要存储的「简单」情况也是如此(比如,只是用 torch.zeros(2, 2) 划配一个邻接张量时)。

顺便一提,我们感兴趣的不是这种情况,而是有一个分立的存储概念的情况,只是将一个域段定义为有一个基张量支持的张量。这会更加复杂一些,但也有好处:邻接张量可以实现远远更加直接的表示,而没有存储造成的间接麻烦。这样的变化能让 PyTorch 的内部表示方式更接近 Numpy。

我们已经介绍了一些张量的数据布局(有人可能会说,如果你正确地理解了数据表示,其它一切都会自然到位)。但还是有必要简要谈谈如何实现对张量的操作。在最抽象的层面上,当你调用 torch.mm 时,会发生两次调度

第一次调度基于设备类型和张量布局:比如是 CPU 张量还是 CUDA张量,是有步幅的张量还是稀疏的张量。这个调度是动态的:这是一个虚函数(virtual function)调用(这个虚函数调用究竟发生在何处是本演讲后半部分的主题)。

这里需要做一次调度应该是合理的:CPU 矩阵乘法的实现非常不同于 CUDA 的实现。这里是动态调度的原因是这些核(kernel)可能位于不同的库(比如 libcaffe2.so 或 libcaffe2_gpu.so),这样你就别无选择:如果你想进入一个你没有直接依赖的库,你必须通过动态调度抵达那里。

第二次调度是在所涉 dtype 上的调度。这个调度只是一个简单的 switch 语句,针对的是核选择支持的任意 dtype。这里需要调度的原因也很合理:CPU 代码(或 CUDA 代码)是基于 float 实现乘法,这不同于用于 int 的代码。这说明你需要为每种 dtype 都使用不同的核。

专业用户独享

本文为机器之心深度精选内容,专业认证后即可阅读全文
开启专业认证
工程PyTorch深度学习框架
131
相关数据
调度技术

调度在计算机中是分配工作所需资源的方法。资源可以指虚拟的计算资源,如线程、进程或数据流;也可以指硬件资源,如处理器、网络连接或扩展卡。 进行调度工作的程序叫做调度器。调度器通常的实现使得所有计算资源都处于忙碌状态,允许多位用户有效地同时共享系统资源,或达到指定的服务质量。 see planning for more details

导数技术

导数(Derivative)是微积分中的重要基础概念。当函数y=f(x)的自变量x在一点x_0上产生一个增量Δx时,函数输出值的增量Δy与自变量增量Δx的比值在Δx趋于0时的极限a如果存在,a即为在x0处的导数,记作f'(x_0) 或 df(x_0)/dx。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

雅可比矩阵技术

在向量分析中,雅可比矩阵是函数的一阶偏导数以一定方式排列成的矩阵,其行列式称为雅可比行列式。在代数几何中,代数曲线的雅可比行列式表示雅可比簇:伴随该曲线的一个代数群,曲线可以嵌入其中。

PyTorch