全面解读PyTorch内部机制
共 11701字,需浏览 24分钟
· 2021-05-31
点击上方“人工智能与算法学习”,选择加"星标"或“置顶”
重磅干货,第一时间送达
本文转自 深度学习这件小事
斯坦福大学博士生与 Facebook 人工智能研究所研究工程师 Edward Z. Yang 是 PyTorch 开源项目的核心开发者之一。他在 PyTorch 纽约聚会上做了一个有关 PyTorch 内部机制的演讲,本文是该演讲的长文章版本。
![](https://filescdn.proginn.com/f54188ee284c1762672f11dc8c00a085/007d97d831f3608656886b0b48bc0141.webp)
![](https://filescdn.proginn.com/d72b340103b333472dec738944e67f19/b62d14357b1d2b674e551424cd58b04c.webp)
![](https://filescdn.proginn.com/48525bae0d2c4a602db4f20fef01d217/fde1d9ae3689963df2f95262bbf11b6f.webp)
![](https://filescdn.proginn.com/a4b2ddd019f3da28d3ca728a04837d94/9fb2b4551c306873ab42a5c961e7f57c.webp)
![](https://filescdn.proginn.com/bb873d14d3c2867c18d7417e654e8d8a/4e6d509b8980a15f28eb4eaaa860dc7f.webp)
演讲时的提问:如果我取张量的一个域段,我该如何释放底层张量的内存? 答案:你必须制作该域段的一个副本,由此断开其与原始物理内存的连接。你能做的其它事情实际上并不多。另外,如果你很久之前写过 Java,取一个字符串的子字符串也有类似的问题,因为默认不会制作副本,所以子字符串会保留(可能非常大的字符串)。很显然,Java 7u6 将其固定了下来。
![](https://filescdn.proginn.com/199b085bd3fc351bc83abacdc107344e/d0eeaad10f7aaba40665eaf62da911dd.webp)
![](https://filescdn.proginn.com/d3838e568468faa2abb423913ff202b5/5a1494e3b6de541d1c96105e824e8ba1.webp)
顺便一提,我们感兴趣的不是这种情况,而是有一个分立的存储概念的情况,只是将一个域段定义为有一个基张量支持的张量。这会更加复杂一些,但也有好处:邻接张量可以实现远远更加直接的表示,而没有存储造成的间接麻烦。这样的变化能让 PyTorch 的内部表示方式更接近 Numpy。
![](https://filescdn.proginn.com/d9ef8fdf063455dba370147fccfa12bc/ce591ccf4d6b0bd232143d483db39545.webp)
![](https://filescdn.proginn.com/e07f3323d98a3534f6bc4aabfa5887e9/2d09f02f0cbd139c139ac0fb116f976c.webp)
![](https://filescdn.proginn.com/6fed04cd2b2c7861111c7687095b38ea/68ad4f58997dc9fb6a1aa811a6422cf9.webp)
device(设备):描述了实际存储张量的物理内存,比如在 CPU、英伟达 GPU(cuda)、AMD GPU(hip)或 TPU(xla)上。设备之间各不相同的特性是有各自自己的分配器(allocator),这没法用于其它设备。
layout(布局):描述了对物理内存进行逻辑解读的方式。最常用的布局是有步幅的张量(strided tensor),但稀疏张量的布局不同,其涉及到一对张量,一个用于索引,一个用于数据;MKL-DNN 张量的布局更加奇特,比如 blocked layout,仅用步幅不能表示它。
dtype(数据类型):描述了张量中每个元素实际存储的数据的类型,比如可以是浮点数、整型数或量化的整型数。
![](https://filescdn.proginn.com/eddf00d94453f59b1ac9529e3a82e66b/55fc27305993eec8a5bfaa5a934f0913.webp)
![](https://filescdn.proginn.com/ac1d79627a1953c83a2ff2174123ffb9/c1d39123a6369f463258364dd32bc9dc.webp)
![](https://filescdn.proginn.com/1338a0f86bbab9b7357b1158ae4dad23/d623d648de13a042b9fffbe6b9a2b86a.webp)
首先将你的目光投向红色和蓝色的变量。PyTorch 实现了反向模式自动微分,这意味着我们可以「反向」走过前向计算来有效地计算梯度。查看变量名就能看到这一点:在红色部分的底部,我们计算的是损失(loss);然后在这个程序的蓝色部分,我们所做的第一件事是计算 grad_loss。loss 根据 next_h2 计算,这样我们可以计算出 grad_next_h2。从技术上讲,我们加了 grad_ 的变量其实并不是梯度,它们实际上左乘了一个向量的雅可比矩阵,但在 PyTorch 中,我们就称之为 grad,基本上所有人都知道这是什么意思。
如果代码的结构保持一样,而行为没有保持一样:来自前向的每一行都被替换为一个不同的计算,其代表了前向运算的导数。举个例子,tanh 运算被转译成了 tanh_backward 运算(这两行用图左边一条灰线连接)。前向和反向运算的输入和输出交换:如果前向运算得到 next_h2,反向运算就以 grad_next_h2 为输入。
![](https://filescdn.proginn.com/92d944c7a4ba22f22485687b1b060660/73a75a5ede902ce58cc5fd6657176ff3.webp)
![](https://filescdn.proginn.com/31894cceb1b208c0a2f17cd44522584d/9640687b726de200a42374a9bd85a0c0.webp)
![](https://filescdn.proginn.com/7e9cf5e314d39ec080aaa022d9b607d6/67de32ea45ec26e9196ad2f9bb216e55.webp)
![](https://filescdn.proginn.com/6feecdad751b73a7fd9db884310141d5/0ba592fb0e4c5c5bce8a447ef6260a8a.webp)
![](https://filescdn.proginn.com/2a886c6b1d172d6351818072f3f1b36a/35cf52a396524b8a91143c5c69e29712.webp)
![](https://filescdn.proginn.com/5c27c90517e3488eb2344d9bca03b105/0506aafb7bd95cbe5bfde10b9cdd72eb.webp)
![](https://filescdn.proginn.com/5cb949867b2a45f7716d3a31476f1f04/ac487e74e2dca8348dc29b04ab0b16b6.webp)
![](https://filescdn.proginn.com/0df06fd1a0986e0cc35618d7dfdf3bbc/ac487e74e2dca8348dc29b04ab0b16b6.webp)
![](https://filescdn.proginn.com/a4af74ec7d28d3d1cd8d0f9488b8491f/39be84f59d7137ab28f3a4b1eecbc127.webp)
![](https://filescdn.proginn.com/a4dfe6d2a43c11604ad44cfe47ab18f7/157934598ec295a6e1c3932c36365538.webp)
首先,torch/ 包含你最熟悉的东西:你导入和使用的实际的 Python 模块。这些东西是 Python 代码而且易于操作(只需要进行修改然后查看结果即可)。但是,如果太过深入……
torch/csrc/:实现了你可能称为 PyTorch 前端的 C++ 代码。用更描述性的术语讲,它实现了在 Python 和 C++ 间转换的绑定代码(binding code);另外还有一些相当重要的 PyTorch 部分,比如 autograd 引擎和 JIT 编译器。它也包含 C++ 前端代码。
aten/:这是「A Tensor Library」的缩写(由 Zachary DeVito 命名),是一个实现张量运算的 C++ 库。如果你检查某些核代码所处的位置,很可能就在 ATen。ATen 本身就分为两个算子区域:「原生」算子(算子的现代的 C++ 实现)和「传统」算子(TH、THC、THNN、THCUNN),这些是遗留的 C 实现。传统的算子是其中糟糕的部分;如果可以,请勿在上面耗费太多时间。
c10/:这是「Caffe2」和「A"Ten"」的双关语,包含 PyTorch 的核心抽象,包括张量和存储数据结构的实际实现。
![](https://filescdn.proginn.com/8d246912d3003080e66d5486db773fd9/1d966c17974f63cc7768ce21b60d4031.webp)
我们必须从 Python 国度转换到 C++ 国度(Python 参数解析)。
我们处理变量调度(VariableType—Type,顺便一提,和编程语言类型并无特别关联,只是一个用于执行调度的小工具)。
我们处理设备类型/布局调度(Type)。
我们有实际的核,这要么是一个现代的原生函数,要么是传统的 TH 函数。
![](https://filescdn.proginn.com/9f52c573fac699160dee74164ef9a202/47408e74d7edb1e1fafebc28dd5d7103.webp)
![](https://filescdn.proginn.com/866fb21300ee2d1d8461d6c976309dc4/4712fada26c0fff8b0d159c35058414f.webp)
![](https://filescdn.proginn.com/bc0a7f4f296c0f31d2412375b62d27fd/67370800ac872b511c301dd14d478848.webp)
![](https://filescdn.proginn.com/de33b29538a9b0588575dcde7b35ac0b/6931f8a0f70c64bd4fd3b92e7de52314.webp)
首先有一些我们要写的有关核的元数据,这能助力代码生成并让你获取所有与 Python 的捆绑包,同时无需写任何一行代码。
一旦你到达了核,你就经过了设备类型/布局调度。你首先需要写的是错误检查,以确保输入的张量有正确的维度。(错误检查真正很重要!不要吝惜它!)
接下来,我们一般必须分配我们将要写入输出的结果张量。
该到写核的时候了。现在你应该做第二次 dtype 调度,以跳至其所操作的每个 dtype 特定的核。(你不应该过早做这件事,因为那样的话你就会毫无用处地复制在任何情况下看起来都一样的代码。)
大多数高性能核都需要某种形式的并行化,这样就能利用多 CPU 系统了。(CUDA 核是「隐式」并行化的,因为它们的编程模型构建于大规模并行化之上。)
最后,你需要读取数据并执行你想做的计算!
![](https://filescdn.proginn.com/6a20599788474eb445ff29a03a2426ad/ea5300729320d11fabe62ff7c3865c3c.webp)
![](https://filescdn.proginn.com/fd71196d3647ea14c83fda64aa0e7065/3f5a7b6d3d9633624078f8cd6c922347.webp)
![](https://filescdn.proginn.com/bd1b7d757e0a41f92c227ecc208d43e0/799729230ab8ed50dc650c9cf8ffadf2.webp)
![](https://filescdn.proginn.com/ca595683e2e909de978819180290715b/bd3d4c61d9d4b0e5cb819157ea1587ab.webp)
![](https://filescdn.proginn.com/3846c488f032260fe0b361dfd230f4d9/4596902260f6caeab03a3953707f3a43.webp)
![](https://filescdn.proginn.com/4cf155e6c7898e9f7b3d55a9317af3d4/79aeb143b1fbd219995d75578b42af0b.webp)
![](https://filescdn.proginn.com/7e8f17bb45af0b653e4eb4f1f021d4de/fd4bc1d20d9e0b9b288f4cafc4466ca9.webp)
如果你只想获取某个特定位置的值,你应该使用 TensorAccessor。张量存取器就像是一个张量,但它将张量的维度和 dtype 硬编码为了模板参数。当你检索一个存取器时,比如 x.accessor
();,我们会做一次运行时间测试以确保张量确实是这种格式;但那之后,每次存取都不会被检查。张量存取器能正确地处理步幅,因此你最好使用它们,而不是原始的指针访问(不幸的是,很多传统的核是这样做的)。另外还有 PackedTensorAccessor,这特别适用于通过 CUDA launch 发送存取器,这样你就能从你的 CUDA 核内部获取存取器。(一个值得一提的问题:TensorAccessor 默认是 64 位索引,这比 CUDA 中的 32 位索引要慢得多!) 如果你在用很常规的元素存取编写某种算子,比如逐点运算,那么使用远远更高级的抽象要好得多,比如 TensorIterator。这个辅助类能为你自动处理广播和类型提升(type promotion),相当好用。
要在 CPU 上获得真正的速度,你可能需要使用向量化的 CPU 指令编写你的核。我们也有用于这方面的辅助函数!Vec256 类表示一种标量向量,并提供了一些能在它们上一次性执行向量化运算的方法。然后 binary_kernel_vec 等辅助函数能让你轻松地运行向量化运算,然后结束那些没法用普通的旧指令很好地转换成向量指令的东西。这里的基础设施还能在不同指令集下多次编译你的核,然后在运行时间测试你的 CPU 支持什么指令,再在这些情况中使用最佳的核。
![](https://filescdn.proginn.com/087642be3330572f95efec04db5ebda8/62211b3c32c873ee8395dc2c998ff7e0.webp)
它是以 C 风格书写的,没有(或很少)使用 C++。
其 refcounted 是人工的(使用了对 THTensor_free 的人工调用以降低你使用张量结束时的 refcounts)。
其位于 generic/ 目录,这意味着我们实际上要编译这个文件很多次,但要使用不同的 #define scalar_t
![](https://filescdn.proginn.com/01e52d3d3f01ad597bff436cad8e8745/9d94feb3200185290692d6920ce28ffe.webp)
如果你编辑一个 header,尤其是被许多源文件包含的 header(尤其当被 CUDA 文件包含时),可以预见会有很长的重新 build 时间。尽量只编辑 cpp 文件,编辑 header 要审慎!
我们的 CI 是一种非常好的零设置的测试修改是否有效的方法。但在获得返回信号之前你可能需要等上一两个小时。如果你在进行一种将需要大量实验的改变,那就花点时间设置一个本地开发环境。类似地,如果你在特定的 CI 配置上遇到了困难的 debug 问题,就在本地设置它。你可以将 Docker 镜像下载到本地并运行:https://github.com/pytorch/ossci-job-dsl
贡献指南解释了如何设置 ccache:https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md#use-ccache ;强烈建议这个,因为这可以让你在编辑 header 时幸运地避免大量重新编译。当我们在不应该重新编译文件时重新编译时,这也能帮你覆盖我们的 build 系统的漏洞。
最后,我们会有大量 C++ 代码。如果你是在一台有 CPU 和 RAM 的强大服务器上 build,那么会有很愉快的体验。特别要说明,我不建议在笔记本电脑上执行 CUDA build。build CUDA 非常非常慢,而笔记本电脑往往性能不足,不足以快速完成。
![](https://filescdn.proginn.com/e64b3550cdcea7c3c2b5e4996d87eb6e/792f01904268982f6ac189ec9114fe9a.webp)