1. 从一张普通显卡说起：为什么我要自己写神经网络框架
三年前我还在用别人封装好的高层API搭模型，调包调得飞起，但每次遇到显存爆炸、梯度异常、训练速度上不去的问题，就只能靠搜索引擎和试错硬扛。直到有一次，我在一台只有8GB显存的普通游戏显卡上想跑一个稍微像样的图像分类任务，结果主流框架的默认配置直接给我报了个显存不足。那一刻我意识到，如果我不懂底层的内存分配逻辑、不懂计算图的构建方式、不懂算子是怎么在GPU上调度执行的，我就永远只能被框架的默认行为牵着鼻子走。
于是我开始动手写一个属于自己的轻量级神经网络框架。目标很明确：能在普通消费级显卡上完成训练，不追求支持超大规模模型，但要把前向传播、反向传播、参数更新、内存管理这些核心环节全部打通。这个项目后来被我开源出去，陆续收到了一些反馈，有人用它跑通了手写数字识别，有人拿它做简单的时序预测，也有人纯粹拿来学习自动求导的原理。
这篇文章不是一篇“从零教你写框架”的教科书式教程，而是把我在这条路上踩过的坑、做过的取舍、以及那些真正让普通显卡跑得动模型的关键设计，原原本本地讲清楚。如果你也曾经被显存问题困扰，或者想搞明白一个神经网络框架到底在背后做了什么，那接下来的内容应该会对你有用。
2. 为什么普通显卡跑不动大框架：显存占用的真实构成
2.1 显存都去哪儿了：拆解训练时的内存开销
很多人以为显存只用来存模型参数，实际上参数本身占的比例可能小得可怜。以一个简单的卷积网络为例，假设参数量是500万，用float32存储也就20MB左右。但训练时显存占用动辄几个GB，多出来的部分主要来自这几个方面：
- 中间激活值：前向传播过程中每一层的输出都需要保留，因为反向传播要用它们计算梯度。这部分占用和批量大小、特征图尺寸直接相关，往往是大头。
- 梯度存储：每个可训练参数都需要一份对应的梯度，占用和参数量相同。
- 优化器状态：如果用的是Adam这类自适应优化器，每个参数还需要保存一阶矩和二阶矩，相当于额外两份参数量的内存。
- 临时缓冲区：矩阵乘法、卷积等操作在GPU上执行时需要临时工作空间，cuDNN之类的库会根据算法选择动态分配。
- 碎片化浪费：频繁申请和释放不同大小的显存块会导致碎片，实际可用显存可能远小于理论值。
我在早期版本中没有做任何内存优化，结果一个参数量不到1000万的小模型，批量大小设为32就爆显存了。后来用工具逐层统计才发现，中间激活值占了总显存的70%以上。
2.2 普通显卡的硬约束：带宽比容量更致命
消费级显卡和高端计算卡之间的差距，很多人只关注显存容量，比如8GB对比40GB。但实际训练时，显存带宽往往是更早遇到的瓶颈。以某款主流消费级显卡为例，显存带宽大约在400-500GB/s量级，而高端计算卡可以到1.5TB/s以上。这意味着同样一次参数更新，普通显卡需要花三倍以上的时间把数据搬来搬去。
这个约束直接影响了框架的设计方向。如果你的框架频繁在CPU和GPU之间同步数据，或者算子实现时没有考虑内存访问的连续性，那训练速度会慢到无法接受。我在设计自己的框架时，把“减少主机与设备之间的数据传输”作为一条硬性原则，所有能放在GPU上完成的操作绝不拿回CPU。
2.3 小批量与梯度累积：用时间换空间的经典策略
既然显存容量有限，最直接的办法就是减小批量大小。但批量太小会导致梯度估计噪声大、训练不稳定。这时候梯度累积就是一个非常实用的技巧：用小的批量多次前向和反向传播，把梯度累加起来，等累积到相当于一个大批量的效果后再更新参数。
我在框架里内置了梯度累积的支持，实现上并不复杂，关键是要注意两点：一是每次累积前不清零梯度，二是累积步数达到后手动触发优化器更新并清零。这个策略让我的8GB显卡能够模拟出批量大小128甚至256的训练效果，代价只是训练时间线性增加。对于不追求极致训练速度的场景，这是性价比最高的方案。
3. 自研框架的核心模块：从张量到自动求导的完整链路
3.1 张量设计：不只是多维数组
张量是框架最基础的数据结构。很多人觉得张量就是带形状的数组，但真正实现起来要考虑的东西远不止于此。我的设计里，一个张量对象需要包含以下信息：
- 数据存储：实际的内存块，可以是CPU内存也可以是GPU显存。
- 形状与步长：支持视图操作，比如转置、切片时不复制数据，只改变步长信息。
- 数据类型：float32、float16、int32等，不同精度对显存和计算速度影响很大。
- 设备标识：标记数据在哪个设备上，方便框架自动调度。
- 梯度相关字段：是否需要计算梯度、梯度值存在哪里、被哪些操作依赖。
这里有一个容易忽略的细节：内存对齐。GPU对内存访问有对齐要求，如果张量的起始地址没有对齐到合适的边界，某些算子会直接报错或者性能急剧下降。我在分配内存时统一按256字节对齐，虽然会浪费一点空间，但避免了大量潜在的兼容性问题。
另一个关键设计是延迟计算。有些操作（比如reshape、transpose）不需要立即执行数据搬运，只需要修改元信息即可。把这些操作做成惰性的，能显著减少不必要的显存拷贝。
3.2 计算图的构建方式：动态图为什么更适合自研
主流框架有动态图和静态图两种模式。静态图需要先定义再运行，优点是能提前优化，缺点是调试困难。动态图则是边执行边建图，写起来像普通Python代码，调试直观。
对于自研框架来说，动态图是更务实的选择。原因很简单：静态图的图优化、内存规划、算子融合这些模块工作量巨大，没有几个人年的投入根本做不好。而动态图虽然单步开销略大，但实现起来清晰得多，也更容易定位问题。
我的实现方式是：每个张量上记录一个grad_fn，指向创建它的操作节点。反向传播时从损失值出发，沿着grad_fn链回溯，依次调用每个节点的反向计算函数。这个链条就是计算图，只不过是在前向过程中动态建立的。
3.3 自动求导的实现细节：反向传播不只是链式法则
自动求导的核心是链式法则，但工程实现上有几个坑：
第一个坑是梯度累加。当一个张量被多个后续操作使用时，它的梯度需要从多个路径累加。比如y = x * 2 + x * 3，x的梯度应该是5而不是2或3。实现时需要在反向传播前把梯度缓冲区清零，然后每个路径贡献的梯度累加上去。
第二个坑是原地操作。像x += 1这种原地修改会破坏计算图的历史信息，导致反向传播时拿不到原始值。我的做法是禁止对需要梯度的张量做原地操作，或者在原地操作前保存一份副本。
第三个坑是梯度检查。自己写的反向传播公式很容易出错，尤其是涉及矩阵乘法、卷积、池化这些复杂算子时。我强烈建议实现一个数值梯度检查工具，用有限差分法计算近似梯度，和解析梯度对比。虽然慢，但能发现绝大多数实现错误。
# 数值梯度检查的简化示例 def numerical_gradient(f, x, eps=1e-4): grad = np.zeros_like(x) for i in range(x.size): x_flat = x.flatten() x_flat[i] += eps f_plus = f(x_flat.reshape(x.shape)) x_flat[i] -= 2 * eps f_minus = f(x_flat.reshape(x.shape)) grad_flat = grad.flatten() grad_flat[i] = (f_plus - f_minus) / (2 * eps) return grad这段代码在CPU上用numpy就能跑，每次检查几个随机位置即可，不需要全量检查。
4. 让普通显卡跑得动的关键优化：内存与计算的双重瘦身
4.1 混合精度训练：用float16换速度和显存
混合精度是让普通显卡跑得动模型的最有效手段之一。核心思路是：前向和反向传播用float16存储激活值和梯度，参数更新时用float32保持精度。这样显存占用直接减半，而且很多GPU对float16有专门的计算加速。
但混合精度不是简单地把所有数据转成float16就完事了。有几个关键点：
- 损失缩放：float16的动态范围小，梯度值容易下溢变成0。需要把损失值放大一个系数，反向传播后再把梯度缩回去。
- 主权重副本：优化器里保存一份float32的参数副本，更新时在float32上做，然后转成float16给下一轮前向用。
- 算子兼容性：不是所有算子都支持float16，比如某些归约操作需要float32累加。框架需要自动处理这些类型转换。
我在框架里实现了一个简单的自动混合精度模块，用户只需要在训练循环里加两行代码就能启用。实测下来，显存占用降低了约40%，训练速度提升了20%左右，精度损失在可接受范围内。
4.2 激活值重计算：时间换空间的极致玩法
激活值重计算（也叫梯度检查点）是另一个显存优化利器。原理很简单：前向传播时只保存部分层的激活值，反向传播需要用到被丢弃的激活值时，重新计算一遍。
这个策略的代价是额外的计算量，通常会增加20%-30%的训练时间。但对于显存极度受限的场景，这是能让你跑起更大模型的关键。我在框架里把重计算做成了可配置的：用户可以选择对哪些层启用，或者设置一个显存阈值，超过阈值自动启用。
实现上需要注意的是，重计算的前向过程必须在反向传播时能够复现，所以随机性操作（如dropout）需要保存随机种子或者使用确定性实现。
4.3 内存池与碎片整理：减少无效占用
频繁申请释放显存会导致碎片化，最终明明有足够的总空闲显存却分配不出连续的大块。解决办法是实现一个内存池：框架启动时预先申请一大块显存，之后所有张量分配都从池子里切分，释放时归还给池子而不是还给系统。
内存池的另一个好处是减少分配开销。GPU显存分配是相对昂贵的操作，从池子里切分几乎不花时间。我的实现里用了简单的首次适应算法，配合定期碎片整理（把空闲块合并），实测能减少约15%的显存浪费。
5. 训练循环的工程细节：那些文档里不会写的东西
5.1 数据加载：别让CPU成为瓶颈
很多人把注意力全放在GPU计算上，结果发现GPU利用率只有30%，瓶颈在数据加载。普通显卡虽然算力有限，但如果数据管道没做好，照样喂不饱。
我的经验是：数据预处理尽量放在CPU上并行做，用多进程而不是多线程。Python的GIL会让多线程在CPU密集任务上几乎没提升，而多进程能真正利用多核。具体做法是维护一个预取队列，后台进程持续把处理好的数据放进队列，GPU需要时直接从队列取。
另外，数据增强操作要尽量轻量。我见过有人在训练循环里做复杂的图像变换，结果数据加载时间比前向反向加起来还长。如果增强逻辑复杂，考虑离线预处理或者用更高效的实现。
5.2 学习率调度：小批量训练更需要精细调参
小批量训练时梯度噪声大，学习率需要更精细的调度。我常用的策略是预热加余弦退火：前几个epoch从很小的学习率线性增加到初始值，然后按余弦曲线慢慢降到接近零。
预热的作用是让模型在训练初期不要被大梯度带偏，尤其是用了批归一化的时候。余弦退火则能让模型在训练后期稳定收敛，避免学习率过大导致在最优解附近震荡。
在框架里实现调度器时，要注意学习率的更新时机。我习惯在每个epoch结束后更新，而不是每个batch，这样调度曲线更平滑。当然如果数据集很大，每个epoch包含很多batch，也可以按步更新。
5.3 梯度裁剪与异常检测：训练稳定性的保险丝
梯度爆炸是训练深层网络时的常见问题，表现是损失突然变成NaN。梯度裁剪是最直接的解决办法：计算所有参数梯度的范数，如果超过阈值就按比例缩放。
但裁剪阈值设多少合适？我的经验是从1.0开始试，如果发现梯度经常被裁剪，说明模型可能太深或者学习率太大。另外，除了裁剪，还应该加一个异常检测：如果损失变成NaN或Inf，立即停止训练并保存现场，方便排查。
我在框架里内置了一个简单的异常检测钩子，每次参数更新前检查梯度是否包含非有限值。这个功能帮我省了很多调试时间。
6. 实测效果与适用边界：这个框架能做什么、不能做什么
6.1 在普通显卡上的实测数据
我用一台搭载8GB显存消费级显卡的机器做了几组测试，对比自研框架和主流框架在相同任务上的表现：
| 任务 | 模型规模 | 批量大小 | 自研框架显存 | 主流框架显存 | 自研框架速度 |
|---|---|---|---|---|---|
| 手写数字分类 | 约120万参数 | 128 | 1.8GB | 2.4GB | 0.85x |
| 小型图像分类 | 约500万参数 | 64 | 3.2GB | 4.5GB | 0.78x |
| 简单时序预测 | 约80万参数 | 256 | 1.2GB | 1.6GB | 0.92x |
速度上自研框架确实慢一些，主要差距在算子优化程度。但显存占用有明显优势，平均低25%-30%。这意味着在同样的硬件上，自研框架能跑更大的批量或者更大的模型。
6.2 什么时候该用自研框架，什么时候不该
自研框架适合这些场景：
- 学习目的：想彻底搞懂神经网络训练的全流程。
- 显存极度受限：主流框架跑不起来的模型，自研框架可能通过更激进的内存优化跑起来。
- 需要深度定制：比如实现特殊的梯度计算方式、自定义算子、非标准训练流程。
- 部署环境特殊：需要在没有主流框架支持的平台上运行。
不适合的场景也很明确：
- 追求极致训练速度：主流框架的算子优化是多年积累的结果，自研很难追上。
- 需要大规模分布式训练：自研框架做分布式的工作量是另一个量级。
- 生产环境要求高稳定性：自研框架的测试覆盖和边界处理肯定不如成熟项目。
我的建议是：把自研框架当作学习工具和特定场景的补充，而不是主流框架的替代品。理解底层原理之后，再用主流框架时你会更清楚每个参数背后的含义，遇到问题也能更快定位。
6.3 后续可以继续深挖的方向
这个项目还有很多可以完善的地方。短期计划是加入对更多算子的支持，比如不同类型的卷积变体、注意力机制相关的操作。中期想尝试实现一个简单的图优化模块，把相邻的逐元素操作融合成一个kernel，减少显存访问次数。长期来看，如果精力允许，想探索一下在边缘设备上的部署，把训练好的模型转成更紧凑的格式。
不过这些都是后话。对于现在这个版本，我最满意的地方是它足够简单，核心代码加起来不到三千行，一个周末就能通读一遍。这种透明性本身就是价值——你知道每一行代码在做什么，出了问题也能自己修。
7. 写在最后：一些个人体会
做这个项目的过程中，最大的收获不是写出了一个能跑的框架，而是对“训练一个神经网络”这件事有了完全不同的理解。以前调包的时候，模型不收敛就换优化器、调学习率、加数据，像在黑暗中摸索。现在我知道每一步背后发生了什么，知道显存里每一MB被谁占用了，知道梯度是怎么从损失函数一步步传回每一层参数的。
这种掌控感让我在面对新问题时更有底气。比如有一次遇到一个奇怪的显存泄漏，我直接在自己的框架里复现了类似场景，用内存池的统计信息定位到是一个中间张量没有被正确释放。如果只用主流框架，这个问题可能要靠猜和试。
如果你也在考虑动手写点什么，我的建议是：不要一上来就追求大而全。先实现一个最小的能跑通线性回归的版本，然后逐步加功能。每加一个功能，确保你完全理解它的原理和实现细节。这个过程会很慢，但每一步都算数。
代码开源之后，有人问我为什么不直接用现成的。我的回答是：用现成的能让你完成任务，但自己写能让你理解任务。这两件事的价值不一样。
