用C++从零实现张量运算与自动微分：拆解PyTorch底层原理

为什么要用C++从零实现深度学习框架

当我们习惯了 import torch 然后调用 loss.backward() 时,很少有人真正思考背后究竟发生了什么。PyTorch 是 Meta（Facebook）开发的开源深度学习框架,于 2016 年发布,以动态计算图（define-by-run）和 Pythonic 接口著称,迅速成为学术界和工业界最流行的深度学习工具之一。然而,其易用性背后隐藏着复杂的底层实现:C++ 核心引擎、CUDA 加速、自动求导系统等。当开发者调用 loss.backward() 时,框架自动完成计算图构建、梯度计算和参数更新,这种高度封装让使用者难以窥见其内部机制。

一位开发者（GitHub 用户 mechanical-turk）决定用 C++ 从零构建一套张量运算库和自动微分引擎,目的正是揭开 PyTorch 的黑箱——看看那些框架在你调用一行代码时到底做了哪些工作。这种"黑箱"特性虽然提升了开发效率,却也让许多开发者对深度学习的本质原理理解停留在表面。

这个名为「Deep Learning All The Way Down」的项目并非为了取代成熟框架,而是一次深度的学习实践。作者不仅公开了完整代码和 Git 检查点,还将整个过程录制成视频系列,目前已更新到第 7 集。这种「边造边讲」的方式,对想深入理解深度学习底层机制的开发者极具参考价值。

C++张量实现:从扁平存储到广播机制

张量（Tensor）是所有深度学习框架的核心数据结构。作者的 C++ 张量目前已经支持了一套相当完整的功能:

• 扁平存储（flat storage）:底层用一维连续数组存放数据,这也是 PyTorch、NumPy 等主流库的通用做法。扁平存储是高性能数值计算库的标准做法。多维数组在逻辑上是二维、三维甚至更高维的结构,但在物理内存中必须以一维连续方式存储。这种设计有三大优势:首先,连续内存布局能最大化 CPU 缓存命中率,现代处理器的缓存预取机制会将相邻数据一并载入,显著提升访问速度;其次,简化了内存分配和释放,只需一次 malloc/free 操作;最后,便于与 BLAS、LAPACK 等底层数学库对接,这些库都假定数据连续存储。

• 多维索引:在扁平存储之上通过步长（stride）计算实现多维坐标到线性偏移的映射。通过步长数组记录每个维度的跨越距离,可以高效实现多维索引、转置、切片等操作,甚至无需复制数据。

• 逐元素运算（elementwise operations）:加减乘除等基础操作。

• 规约操作（reductions）:如求和、求均值等降维运算。

• 广播（broadcasting）:让不同形状的张量能够按规则自动对齐参与运算。广播是 NumPy 在 2005 年引入的概念,后被所有主流深度学习框架继承。它定义了一套规则,让形状不同的数组能够参与算术运算。核心规则是:从最右侧维度开始对齐,若某维度大小为 1 或缺失,则自动"拉伸"以匹配另一数组。例如,形状 (3, 1) 的数组与 (3, 4) 相加时,(3, 1) 会在第二维度上广播为 (3, 4)。这种机制在深度学习中无处不在:批量归一化时减去均值、注意力机制中的掩码相乘、损失函数计算等。实现广播需要复杂的形状推断逻辑和步长调整,这是深度学习中极为常用但实现起来颇为繁琐的机制。

• 二阶矩阵乘法（rank-two matmul）:神经网络中最核心的计算之一。

• 均方误差（MSE）:常见的损失函数实现。

从这份功能清单可以看出,作者已经构建了一个能支撑基本神经网络前向计算的张量系统。这些看似基础的能力,正是任何深度学习框架的地基。

标量自动微分引擎:反向传播的核心实现

项目最近的重头戏,是一套独立的标量反向模式自动微分（scalar reverse-mode autograd）引擎。自动微分（automatic differentiation, AD）有前向模式和反向模式两种。深度学习中普遍使用反向模式,因为它能用 O(n) 的时间计算出一个标量损失对所有参数的梯度,而前向模式需要 O(n) 次前向传播。反向模式基于微积分的链式法则:若 z=f(y) 且 y=g(x),则 dz/dx = (dz/dy) × (dy/dx)。在计算图中,每个操作节点存储其局部导数,反向传播时,从输出节点开始,将上游梯度乘以局部导数,逐层传递回输入节点。它的工作流程清晰地还原了现代自动微分的核心思想。

计算图构建与反向传播流程

1. 前向传播时构建计算图:算术运算符在执行的同时,动态记录下操作之间的依赖关系,形成一张计算图。计算图是深度学习框架的核心抽象:节点代表操作或变量,边代表数据依赖。静态图（TensorFlow 1.x）在执行前完整构建,动态图（PyTorch）边计算边构建。

2. backward() 生成拓扑排序:从输出（损失）节点出发,对整张图进行拓扑排序。反向传播需要按依赖关系逆序访问节点,这正是拓扑排序的应用场景。拓扑排序是有向无环图（DAG）的经典算法,常用 Kahn 算法或深度优先搜索实现。在自动微分中,正确的拓扑排序保证了梯度计算的顺序:子节点的梯度必须在父节点之前计算完毕。

3. 反向遍历图:按拓扑逆序依次访问每个节点。

4. 应用局部导数:每个操作根据链式法则应用自己的局部梯度。

5. 梯度累加:当某个值通过多条路径影响最终损失时,来自不同路径的梯度会被正确地累加。梯度累加在自动微分中至关重要:当一个变量通过多个计算路径影响损失函数时,来自各路径的梯度必须相加。例如,若 x 既直接参与 loss 计算,又通过中间变量 y 间接影响 loss,则 dL/dx = (dL/dx)_direct + (dL/dy) × (dy/dx)。链式法则的加法版本保证了这一点。

最后这一点尤为关键——梯度累加正是自动微分正确性的核心,也是初学者最容易忽视的细节。

梯度计算验证示例

作者给出了一段简洁的示例代码:

Value prediction = w1*x1 + w2*x2 + w3*x3 + bias;
Value residual = prediction - target;
Value loss = residual * residual;
loss.backward();

给定权重 [0.5, -1.0, 2.0]、输入 [4.0, 3.0, 2.0]、偏置 0.5、目标值 2.5:

• 前向传播得到预测值 3.5,损失 1.0
• 反向传播恢复出梯度:dL/db = 2,dL/dw = [8, 6, 4]

我们可以手动验证:残差为 3.5 - 2.5 = 1.0,损失对残差的导数为 2 × 1.0 = 2.0（即 dL/db）;而对每个权重的导数则是 2.0 × 对应输入,即 [8, 6, 4]。结果完全吻合。这种可手算验证的设计,让学习者能建立起对梯度计算的直觉信心。

张量与自动微分整合:悬而未决的设计难题

目前,标量自动微分引擎与张量实现是彼此独立的两套系统。作者的下一步计划,是将图的身份标识（graph identity）、所有权（ownership）和梯度信息整合进张量,从而搭建起完整的训练循环。

为此他在社区抛出了一个颇具工程深度的问题:

对于张量与自动微分的整合,应该把 autograd 元数据放在每个 Tensor 句柄内部,还是让张量指向独立的共享计算图节点?

这实际上触及了深度学习框架设计的一个根本抉择:

• 元数据内置于张量:实现直接,张量与其梯度信息紧耦合,但在张量被复制、共享或视图化时,图的所有权管理会变得复杂。

• 张量指向共享图节点:更接近 PyTorch 的实际设计思路。PyTorch 的 Tensor 通过 AutogradMeta 结构关联到计算图。AutogradMeta 包含指向 Node（图节点）的指针、梯度 Tensor、版本计数器等信息。这种设计实现了张量与计算历史的解耦:多个 Tensor 可以共享同一 Node（如 in-place 操作前后的视图）,Node 的生命周期由引用计数管理,独立于 Tensor。当调用 .detach() 时,新 Tensor 的 AutogradMeta 为空,切断了与计算图的联系。这种方法能更优雅地处理多个张量共享同一计算历史的情况,但引入了引用计数和生命周期管理的复杂性。

PyTorch 本身选择了后者的变体:Tensor 通过 AutogradMeta 关联到计算图,图节点独立于张量存活。对于一个学习性项目而言,理解这一抉择背后的权衡,本身就是极有价值的收获。

对深度学习开发者的启示

这个项目最大的价值,不在于它能否媲美 PyTorch 的性能,而在于它把深度学习框架去神秘化了。通过亲手实现扁平存储、广播、拓扑排序和梯度累加,作者（以及跟随视频系列的学习者）建立起了对框架底层机制的第一性理解。

对于任何想真正吃透深度学习的人来说,「造一个玩具版 PyTorch」都是一条被反复验证的有效路径。micrograd 是 Andrej Karpathy（OpenAI 创始团队成员、前 Tesla AI 负责人）开发的极简自动微分引擎,核心代码不到 150 行 Python。它实现了标量级别的反向传播,支持构建简单神经网络。Karpathy 将其作为教学工具,配套视频课程"从零构建神经网络",帮助学习者理解 autograd 本质。类似项目还有 tinygrad（George Hotz 的作品,约 1000 行实现完整 GPU 加速）等。这些"玩具框架"的价值在于:剥离工程复杂度后,暴露出深度学习的数学核心。通过复现这类项目,开发者能深刻理解动态计算图、链式法则、梯度下降的本质,而不是停留在调用 API 的层面。

而用 C++ 而非 Python 来实现,则额外逼迫开发者直面内存管理、所有权和性能等更硬核的系统问题。

项目代码与 Git 检查点已在 GitHub 开源,配套的视频系列也在持续更新中。对于想动手复刻的开发者,这既是一份可运行的参考实现,也是一个观察「造轮子」思考过程的窗口。

核心要点

• 该项目用 C++ 从零实现张量运算和自动微分,目标是揭开 PyTorch 等框架的内部机制
• 张量实现采用扁平存储、步长索引和广播机制,这些是所有主流框架的核心设计
• 标量自动微分引擎展示了计算图构建、拓扑排序和梯度累加的完整流程
• 项目面临的核心设计问题是如何整合张量与计算图,这反映了真实框架的工程权衡
• 通过"造轮子"建立对深度学习的第一性理解,是被 Karpathy 等专家反复验证的学习路径