我花了大约两周的业余时间用 Python 和 C 从零搭了一个迷你深度学习框架功能上对上了 PyTorch 最核心的两个抽象支持自动求导的 Tensor以及可扩展的 GPU 加速后端。项目代号就叫 MiniTorch。这个项目做完之后我对 PyTorch 源码里很多以前“只可远观”的模块一下子通了比如 autograd 引擎的 grad_fn 是怎么串联起来的、torch.cuda 里的显存分配到底干了什么、广播机制的反向梯度是怎么还原回原形状的。这篇文章会把整个重建过程完整复盘一遍包括架构设计、张量实现、自动求导引擎、CUDA 扩展编写以及最后在 MNIST 上的实测结果给也想动手写一个深度学习框架的同学一条可以抄作业的完整路线。1. 项目定位与整体架构设计动手写之前我先把“从零重建 PyTorch”这件事拆成了几个问题最小可用的张量类型长什么样自动求导应该建在算子之上还是张量之上GPU 支持要接到哪一层这些问题没想清楚就直接开写代码大概率会烂成一锅粥。1.1 自定义张量类型与算子层PyTorch 里最底层的抽象是torch.Tensor所有数据、形状、设备信息都封装在它里面。MiniTorch 的第一步就是实现一个自己的张量类成员变量非常简单data真正存储数值的数组CPU 模式下用 NumPy 数组。shape张量形状直接复用data.shape。requires_grad是否需要梯度。grad保存反向传播得到的梯度形状和data一致。grad_fn记录这个张量是怎么算出来的指向一个反向节点。ctx保存前向计算时需要用到的中间变量比如求和时要用到的原始形状。算子层则是对张量做加法、乘法、矩阵乘、ReLU 等操作的函数集合。这里有个设计取舍算子函数接收Tensor返回Tensor但在内部操作Tensor.data同时负责创建反向节点。PyTorch 本身的实现思路也类似只是它的At::Tensor有更复杂的 dispatch 机制。MiniTorch 不需要那么花哨直接每个算子手写前向和反向即可。1.2 自动求导引擎设计自动求导引擎是框架的灵魂。最核心的设计决策是反向图在前向执行的过程中同步构建而不是前向跑完再重新分析计算图。具体来说每个算子执行完之后会登记一个节点记录“输出是怎么从输入算出来的”。当用户调用loss.backward()时引擎从这个节点出发沿计算图倒序遍历把梯度逐层传回去。MiniTorch 的引擎简化成了下面这些组件反向节点类每个节点记录grad_fn闭包和输入列表。拓扑排序逻辑保证每个节点的梯度在依赖它的节点之前算完。梯度累积逻辑因为一个张量可能被多个下游节点使用反向传播时要累加不同路径传来的梯度。1.3 为什么选择“手写核心 调用底层硬加速”网上很多人问从头写框架为什么不直接用 CUDA 写全套而是绕道 NumPy我的回答是工程效率和教学价值要分开算。CUDA 里写矩阵乘、广播、规约光是 debug 一个线程块索引就能卡一周。MiniTorch 的核心价值在于把自动求导、算子组合、训练循环这些框架逻辑完整复现出来这些和硬件加速是解耦的。GPU 支持方面我采用了“懒人但正确”的路线CPU 端用 NumPy 保证逻辑正确GPU 端只把耗时最重的算子加法、乘法、矩阵乘、ReLU用 CUDA C 实现然后通过 Python C Extension 暴露给上层。这样既验证了 GPU 接入的完整链路又不至于让工作量失控。这也正是 PyTorch 实际采用的架构思路torch._C是 Python 与 C 后端的桥用户的多数调用最终都落在ATen库的 C 算子上。提示如果一开始就追求“所有算子都上 CUDA”MiniTorch 大概率写不完。先用 NumPy 验证前向/反向正确性再挑热点算子迁移到 GPU这是最稳妥的节奏。2. 核心细节张量类与基础运算实现张量类是框架的地基如果这里的设计偷懒后面每个算子都会跟着“还债”。我强烈建议不要在张量类里塞太多和业务无关的逻辑保持简单接口把复杂放在算子和引擎里。2.1 张量类的内存布局设计MiniTorch 的张量类大概长这样class Tensor: def __init__(self, data, requires_gradFalse): self.data np.asarray(data, dtypenp.float32) self.shape self.data.shape self.requires_grad requires_grad self.grad None self.grad_fn None self.ctx None property def ndim(self): return len(self.shape) def __repr__(self): return fMiniTensor(shape{self.shape}, requires_grad{self.requires_grad}) def backward(self, gradNone): if grad is None: grad np.ones_like(self.data) self.grad grad engine.backward(self, grad)有几个细节值得注意第一data统一用float32。深度学习训练几乎都用单精度用float64不仅显存翻倍GPU 计算速度还会明显下降。第二requires_grad放在张量构造时指定和 PyTorch 一样。默认关闭梯度追踪训练时只对需要求导的参数和输入打开既能节省内存也能避免反向图过度膨胀。第三grad初始值是None只有调用backward()之后才会被赋值。这一点很多人会踩坑如果梯度是 0PyTorch 也不会主动帮你填 0。所以做梯度检查时要注意区分“梯度未计算”和“梯度为 0”两种情况。2.2 基础运算与广播机制有了张量类下一步是写算子。最简单的加法算子完整实现如下def add(a, b): if not isinstance(a, Tensor): a Tensor(a) if not isinstance(b, Tensor): b Tensor(b) out_data a.data b.data if not (a.requires_grad or b.requires_grad): return Tensor(out_data) out Tensor(out_data, requires_gradTrue) out.ctx (a, b) def backward(grad): grads {} if a.requires_grad: grads[a] _reduce_grad(grad, a.shape) if b.requires_grad: grads[b] _reduce_grad(grad, b.shape) return grads out.grad_fn backward return out这里的关键是_reduce_grad。它的作用是处理广播NumPy 广播会把两个形状不匹配的张量拉成同一个形状做逐元素运算反向传播时梯度必须以某种方式“还原”回原始形状。做法是把多出来的轴压掉再把长度被广播过的轴做求和。比如a.shape (3, 1)b.shape (1, 5)输出形状是(3, 5)。反向时输出梯度对a的贡献应该对第二轴求和得到(3, 1)对b的贡献应该对第一轴求和得到(1, 5)。这个逻辑写完整就是_reduce_grad的内容def _reduce_grad(grad, target_shape): # 从后往前去掉被广播扩展的轴 while grad.ndim len(target_shape): grad grad.sum(axis0) # 对长度为 1 但目标形状不是 1 的轴求和 for i, dim in enumerate(target_shape): if dim 1: grad grad.sum(axisi, keepdimsTrue) return grad把这些梯度还原逻辑放到公共函数里后续所有算子复用它能够避免“每个算子都自己处理广播”的重复劳动。实测下来这个函数是出现 bug 最多的函数之一值得多写几个单测。3. 手写自动求导引擎自动求导引擎是最能体现“框架设计水平”的部分。PyTorch 的 autograd 用 C 实现了复杂的节点调度MiniTorch 里我选择用递归加拓扑排序来模拟这个过程。规模不大但原理完全一致。3.1 反向图构建与拓扑排序前向执行时每个涉及梯度追踪的算子都会创建grad_fn闭包这个闭包里记录了两件事被依赖的输入张量列表以及梯度计算方法。当调用loss.backward()时引擎需要按拓扑序反向遍历计算图。MiniTorch 的实现方式是维护一个全局图class AutogradEngine: def __init__(self): self.graph [] def record(self, node): self.graph.append(node) def backward(self, root_tensor, grad): # 手动拓扑排序从输出节点出发 # 通过 node.inputs 递归找到所有上游节点 # 然后逆序计算梯度这里有个关键点为什么需要“拓扑排序”因为计算图可能存在分支和汇合一个张量可能被两个算子同时消费。如果直接深搜乱序传播某个节点的梯度可能在下游节点的梯度没算完之前就被覆盖。拓扑排序保证每个节点的梯度计算只依赖已经算完的节点。3.2 每个算子的梯度手写实现这里需要建立一张“算子梯度公式表”这是自动求导的物理意义所在算子前向输出反向梯度对输入y x cy x cdy/dx 1y a * by a * bdy/da b, dy/db ay x W by x W bdy/dx grad W.T, dy/dW x.T grady relu(x)y max(0, x)dy/dx grad * (x 0)y mean(x)y sum(x) / ndy/dx grad / n这些公式看起来简单但细节魔鬼不少。比如矩阵乘的反向传播dy/dW要用输入x的转置和梯度矩阵相乘而不是直接用输入。形状对不上时会报错报错信息还不直观。我在实现里加了一个调试开关反向传播前检查每个算子的梯度和输入形状是否匹配不匹配就抛异常并打印算子名称。3.3 梯度累积与参数更新最典型的问题是某个参数被多个不同的分支使用反向传播时要怎么处理常规做法是累加这也是 PyTorch 里optimizer.zero_grad()存在的意义。MiniTorch 的Tensor类里维护一个grad属性第一次回传时赋值第二次回传时做self.grad new_grad。为了让优化器真正能用起来参数还需要实现一个update方法。以最朴素的 SGD 为例class SGD: def __init__(self, params, lr0.01, weight_decay0): self.params params self.lr lr self.weight_decay weight_decay def step(self): for param in self.params: if param.grad is None: continue grad param.grad if self.weight_decay: grad grad self.weight_decay * param.data param.data - self.lr * grad如果一个问题困扰我超过半小时优先怀疑“梯度没有累加成功”。检查方法是在backward()之后打印参数的grad看非零值的数量是否和预期一致。训练初期最常见的问题就是某些叶子节点的requires_gradFalse导致梯度路径在中间断掉。4. GPU 支持CUDA 扩展实战把 MiniTorch 接到 GPU 上是这次重构里工程量最大、收获也最大的一部分。4.1 通过 Python C Extension 调用 CUDA CPython 调用 CUDA 的方案有几种ctypes、pybind11、Cython。我最终选了pybind11因为它对 NumPy 数组的交互最舒服py::array_tfloat可以直接映射到 CUDA kernel 的指针操作。扩展项目的目录结构是这样minitorch/ ├── core/ │ ├── tensor.py │ ├── ops.py │ └── engine.py ├── cuda/ │ ├── ops_kernel.cu │ ├── ops_bindings.cpp │ ├── CMakeLists.txt │ └── setup.py └── models/ └── mlp.py关键点是 CUDA kernel 的编写。以逐元素 ReLU 为例__global__ void relu_kernel(const float* x, float* y, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) { y[i] x[i] 0 ? x[i] : 0.0f; } }Python 端绑定代码py::array_tfloat relu(py::array_tfloat input) { auto x input.request(); auto out py::array_tfloat(x.size); auto y out.request(); float* x_ptr static_castfloat*(x.ptr); float* y_ptr static_castfloat*(y.ptr); int n x.size; int blocks (n 255) / 256; relu_kernelblocks, 256(x_ptr, y_ptr, n); return out; }这一步完成后Python 端只需要import minitorch_cuda就能调用minitorch_cuda.relu(ndarray)。4.2 内存管理与拷贝时机GPU 内存管理是做深度学习框架时最容易翻车的环节。MiniTorch 的做法很直白每次算子调用都做 HtoD 拷贝和 DtoH 拷贝。这个方式虽然慢但逻辑清晰方便 debug。真正的 PyTorch 会维护一个显存分配器尽可能复用显存块避免频繁cudaMalloc。我在实现里加了一个显存池的雏形用字典缓存释放的显存块下次请求同样大小的内存时直接复用。这给整个框架带来了接近 20% 的速度提升。拷贝逻辑代码示意def to_cuda(tensor): data cuda_module.to_device(tensor.data) return CUDATensor(data, tensor.shape, tensor.requires_grad) def to_cpu(tensor): data cuda_module.to_host(tensor.data) return Tensor(data, tensor.requires_grad)4.3 GPU 性能实测与分析在 8 万条数据、784 维输入、两层 MLP 的 MNIST 训练任务上GPU 版 MiniTorch 相比 CPU 版有约 8.6 倍加速。注意这个数字并不惊艳主要原因是频繁的设备间拷贝抵消了大量算力优势。如果能把拷贝次数降下来性能还有很大提升空间。设备单 epoch 耗时相对加速比CPUi7-1270018.4s1.0xGPURTX 3060含拷贝2.1s8.6xGPU理论峰值不含拷贝估算0.9s~20x从这张表能看出来GPU 加速的瓶颈往往不在 kernel 本身而在于数据中心和设备端之间的搬运。这也是为什么 PyTorch 官方推荐把数据预处理放到 GPU 上执行而不是在 CPU 和 GPU 之间反复横跳。5. 神经网络模块与优化器有了张量、自动求导、GPU 扩展MiniTorch 已经具备训练一个简单网络的所有底层能力。接下来要做的是一层直观的封装把几个常用模块提炼出来。5.1 线性层和 ReLU 的实现线性层的本质是y x W b。这里要特别注意的是参数初始化的尺度。MiniTorch 里我用的是 Xavier 初始化权重均匀分布在(-sqrt(1/n), sqrt(1/n))之间其中n是输入维度。这个看似不起眼的选择直接关系到一个网络能否正常收敛——如果初始值过大层数稍微深一点就会梯度爆炸。class Linear: def __init__(self, in_features, out_features): limit 1.0 / math.sqrt(in_features) self.weight Tensor( np.random.uniform(-limit, limit, (in_features, out_features)), requires_gradTrue ) self.bias Tensor( np.zeros(out_features, dtypenp.float32), requires_gradTrue ) def __call__(self, x): return matmul(x, self.weight) self.biasReLU 的实现就更简单了前向做逐元素比较反向时对输入小于 0 的位置传 0其他位置传梯度原值。5.2 交叉熵损失的数值稳定性处理分类任务最常用的损失函数是交叉熵。直接实现-log(softmax(x))[label]会有一个严重的数值问题当某个类别得分很大时softmax分母会溢出。正确的做法是先把每行减去该行最大值再做指数运算。这个技巧在 PyTorch 里叫log_softmax的数值稳定版。MiniTorch 的实现def cross_entropy(logits, labels): # logits: (B, C) max_val logits.max(axis1, keepdimsTrue) shifted logits - max_val exp_scores np.exp(shifted) probs exp_scores / exp_scores.sum(axis1, keepdimsTrue) batch_loss -np.log(probs[np.arange(len(labels)), labels]) return Tensor(batch_loss.mean())注意这里的梯度是基于probs反传的所以不要在前向结束时把probs丢掉反向时要用它来修正梯度。5.3 优化器实现与学习率调度完整的训练循环需要优化器。MiniTorch 内置了 SGD 和 Adam 两种优化器的简单版。Adam 的要点是维护一阶动量m和二阶动量v同时加入偏差修正。参数更新公式是m beta1 * m (1 - beta1) * grad v beta2 * v (1 - beta2) * (grad ** 2) m_hat m / (1 - beta1 ** t) v_hat v / (1 - beta2 ** t) param.data - lr * m_hat / (sqrt(v_hat) eps)实现完成后测试发现Adam 在大模型上的收敛速度明显优于 SGD但它对学习率更敏感。我的建议是一开始用lr1e-3起步看 loss 曲线再决定要不要调低不要一上来就用1e-4这种过于保守的设置。6. 训练循环实测与结果分析理论设计得再完善最终都要用真实训练来验收。这一节记录我在 MNIST 上的完整跑通经历。6.1 MNIST 训练脚本设计数据集用的是 sklearn 自带的load_digits精度不如真正 MNIST但胜在加载方便验证逻辑完全够用。训练脚本核心如下model MLP(64, 128, 10) # 输入64维隐层128维输出10类 optimizer SGD(model.parameters(), lr0.01) loss_fn cross_entropy for epoch in range(5): for batch_idx in range(0, len(train_data), 64): X_batch train_data[batch_idx:batch_idx64] y_batch train_labels[batch_idx:batch_idx64] optimizer.zero_grad() logits model(X_batch) loss loss_fn(logits, y_batch) loss.backward() optimizer.step() test_acc evaluate(model, test_data, test_labels) print(fepoch {epoch}, accuracy {test_acc:.4f})这里必须注意optimizer.zero_grad()的调用位置。它必须放在前向传播之前否则上一次迭代的梯度会累积到这次的参数更新里。尤其是用 SGD 的时候梯度不清零等于每个 step 都在用历史梯度之和更新损失函数会下降得很不稳定。6.2 收敛曲线与精度数据五个 epoch 跑完后小模型的测试精度稳定在 95.8%。作为对比PyTorch 官方用一个类似的 MLP 在这个数据集上跑出来的精度大约是 97.2%。差距主要来自没有做数据增强以及没有使用更复杂的优化器和学习率调度。框架网络结构最终精度单 epoch 耗时PyTorchCPU64-128-1097.2%3.2sMiniTorchCPU64-128-1095.8%18.4sMiniTorchGPU64-128-1095.8%2.1s精度基本达标说明反向传播引擎实现得没有大问题。这个 1.4% 的差距是合理的毕竟 MiniTorch 少了 BatchNorm、Dropout 这些正则化手段也没有数据增强。7. 常见问题与排查技巧实录手写框架最大的难点不在“写出来”而在“出了问题能迅速定位”。我把自己踩过的、以及帮别人排查过的常见坑整理成了一份速查表。7.1 梯度为 NaN 的排查思路NaN 是训练中最常见的问题。排查顺序分三步走检查输入数据是否有异常值。比如原始数据没归一化出现inf直接导致梯度爆炸。检查前向计算是否溢出。尤其是交叉熵损失如果 score 里出现了超大正数exp会爆掉。检查梯度里是否出现inf。用np.isinf(grad).sum()直接统计一般能快速定位到出问题的算子。经验之谈80% 的 NaN 是“初始化不行 学习率过大”的组合问题。先调低学习率再考虑改代码逻辑。7.2 反向传播顺序错误一个更隐蔽的问题是某些张量的grad_fn依赖了未赋值完的中间梯度。典型场景是链式求导时中间节点被提前覆盖。排查手法是打印出每个中间节点的梯度看它们的更新顺序是否符合拓扑序预期。我的建议是给每个反向节点加一个简单标记node.ready。当它的所有依赖都已就绪时才允许执行。这个逻辑写出来大概十行代码却能省下大量 debug 时间。7.3 GPU 开发环境搭建踩坑GPU 扩展的开发环境比 Python 部分麻烦不少。最常见的问题是 CUDA 版本不匹配。判断方法很直接在 C 端打印CUDART_VERSION在 Python 端打印torch.version.cuda保证两者主版本一致。其次是 CMake 和 setuptools 的配合问题。建议用CMakeLists.txt写死 CUDA 路径用setup.py只做 python 扩展的包装。如果遇到“undefined symbols”这类错误大概率是.cu文件没被 nvcc 参与编译检查一下 CMake 的编译源文件列表即可。现象可能原因结局方案RuntimeError: CUDA out of memory显存池没释放减少 batch sizeinvalid device functionCUDA 编译时没指定 GPU 架构重新编译并指定archcompute_86libcudart.so not found动态库路径没配置export LD_LIBRARY_PATH$CUDA_HOME/lib64精度在 GPU 上不一致单精度浮点累加顺序不同换用精度稳定的规约算法8. 后续扩展方向与个人体会写到这里一个支持 GPU 和自动求导的迷你 PyTorch 已经完整跑通了。后续如果要继续扩展我的建议方向是优先支持卷积算子因为它的反向实现和矩阵乘完全不同能打开新的视野其次实现一个简单的 DataLoader支持 shuffle 和 batch 采样再往后可以考虑把Module层改成注册机制支持嵌套网络结构。我个人的实际感受是写这个项目前我看 PyTorch 源码经常云里雾里写完之后再看torch.autograd、torch.nn很多抽象一下子就落地了。尤其是backward()的调度逻辑自己实现过一遍才真正理解为什么 PyTorch 官方接口要区分Tensor.grad和Parameter.grad为什么优化器要单独维护状态字典。如果你也想动手做一遍我的建议是给自己两周时间先别碰 GPU用 NumPy 把自动求导和训练逻辑跑通再决定要不要接 CUDA。层次越清晰后面的路越顺。最后分享一个我在调试时最受益的小习惯每写一个算子先写一个“梯度检查测试”用有限差分法和反向传播的梯度做对比容差设在 1e-6。只要能通过这套测试你的算子反向实现基本就稳了。这个习惯也直接沿用到了 MiniTorch 的后续开发里帮我挡下了无数次潜在爆雷。