整个 PyTorch 都在一页上

图 1.整个系列所涉及的程序。您已经输入类似的内容一千次了。本系列的存在是为了让您在结束时了解这些线条的所有作用。所有这些：他们接触的Python、他们使用的C++、他们记录的图表、他们选择的内核、他们使用的内存以及他们运行的两个时钟。这些词中的每一个在其下面的楼层都有一个简单的含义。这是第 0 部分，地图。首先，我们快速地遍历所有层一次。然后我们绘制领土。然后是使代码库的其余部分变得可预测的十二个想法。那么这个系列是如何运作的，以及如何阅读它。这里没有任何内容能讲述完整的故事。这里的一切都有一个地方，每个完整的故事都有一个编号的部分等待着它。在我们开始之前先做出一个承诺。本系列中的每个测量数字都来自一个您可以自己运行的小脚本，该脚本链接在数字出现的位置。我在配备手电筒 2.11.0 的 Apple M3 Max 笔记本电脑上测量了这些结果 [1]。你的数字会有所不同。他们制作的图案不会。 PyTorch 很深。键盘和芯片之间有八个层级。我将它们称为楼层，这个仪表显示了所有楼层。它贯穿整个系列，所以你总是知道自己有多深。图 2. 深度计。橙色点标记您所在的位置。了解一座建筑最快的方法就是不停地走下去一次。就是这个部分。 Floor 2 of 8 python 完整故事：第 5 部分，机械 torch.randn 看起来像一个 Python 函数。询问 Python 它实际上是什么： >>> type ( torch . randn ) Python 仅向用编译代码编写的函数提供该类型。编译代码意味着：在安装之前已将其转换为机器指令的代码，因此其中没有可供读取的 Python 主体，也没有可供调试器停止的行。那么这些机器指令在哪里呢？在共享库中。共享库是程序运行时加载的编译代码文件。它们位于磁盘上的 torch 包内，您可以查看它们（证明）： torch._C -> _C.cpython-312-darwin.so（49 KB，加载程序） libtorch_cpu.dylib 206.5 MB（张量和内核） libtorch_python.dylib 28.5 MB（边界的 python 一侧） p0_the_library.py 证明，准备读取或运行“”证明：pytorch._C 的编译部分实际上是一个精简的编译存根；打印文件及其大小。 _C. __file__ print ( f "torch { torch . __version__ } " ) print ( f "torch._C -> { os .path .basename (stub ) } " f "( { os .path . getsize (stub ) / 1024 :.0f } KB 存根)" ) libdir = os .小路 。 join ( os .path .dirname (stub ), " lib " ) for lib in [ " libtorch_cpu.dylib " , " libtorch_python.dylib " ]: p = os .小路 。加入 ( libdir , lib ) 如果 os .小路 。 contains ( p ): print ( f " { lib :24s } { os .path . getsize ( p ) / 1024 / 1024 :6.1f } MB" ) 下载并运行它 读取大小，然后查看它们： 图 3. 按文件大小按比例绘制。 python可以看到的pytorch部分是橙色的点。您可以从 Python 中看到的 PyTorch 部分是一个 49 KB 文件，其唯一的工作就是加载其他两个文件。真实的主体是 235 MB 的编译代码。 import torch 将其带入您的进程中，之后，调用 torch.randn 意味着跳入该主体。今天我们只需要知道这些文件存在即可。这是代码库的第一个诚实的惊喜：你整天编写的 Python 是它的最小层。 Floor 3 of 8 边界全文：第 5 部分，机械 该调用立即离开 Python。它降落在哪里？在名为 THPVariable_randn 的 C++ 函数中，位于最后一层的 28.5 MB 库内。这里有一个奇怪的事实，你可以保留：这个函数在 PyTorch 存储库中不存在。克隆存储库，搜索名称，但什么也没找到。程序在构建期间编写此函数及其数千个兄弟函数。下面的想法 4 解释了原因，第 5 部分展示了执行写入操作的程序。图 4. 两种语言之间的边界。每个张量运算都会穿过它。跨越这个边界需要时间。要单独查看成本，请计算最小可能操作的时间，其中几乎没有算术隐藏它（证明）：add，1个元素：每次调用 add 0.538 微秒，4M 个元素：每次调用 337.264 微秒 p2_dispatch_cost.py 证明，准备读取或运行“”“证明：一个急切操作的固定成本，以及为什么 size 隐藏它。在两个大小下乘以相同的“a + b”。一个元素的添加几乎是纯机械的（调度、包装、分配）；4M 元素的添加几乎是纯算术，单进程 """ import time import torch def per_op_us ( a , b , iters ): for _ in range ( 2000 ): a + b t0 = time 。 perf_counter () for _ in range ( iters ): a + b return ( time . perf_counter () - t0 ) / iters * 1e6 tiny = per_op_us ( torch .ones ( 1 ), torch .ones ( 1 ), 200_000 ) big_n