1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通了事。我刚开始接触这个领域的时候也是这么想的觉得底层的东西有框架、有库、有现成的轮子何必自己从头造。直到有一次线上推理服务在高峰期出现了一个非常诡异的延迟抖动排查了整整两天最后发现是某个第三方库在特定输入尺寸下触发了隐式的内存重排。那一刻我才意识到如果你对底层的计算图、内存布局、算子调度没有概念你连问题出在哪一层都定位不了。ai-engineering-from-scratch这个标题核心不在于“AI”而在于“from scratch”。它代表的是一种学习路径和工程习惯不满足于当调包侠而是亲手把关键环节拆开、揉碎、再组装一遍。这篇文章适合那些已经会用主流框架跑通模型但总觉得心里没底、想搞清楚“黑盒里面到底发生了什么”的工程师。我会从计算图构建、算子实现、内存管理、推理调度这几个维度把从零搭建一个最小可用AI工程链路的过程讲清楚同时把我在这个过程中踩过的坑和总结的经验一并分享出来。需要提前说明的是从零手搓不等于拒绝使用任何库。合理的做法是核心链路自己实现辅助工具该用就用。你要练的是内功不是重新发明螺丝刀。2. 计算图与自动微分从标量到张量的最小实现2.1 为什么先做标量级自动微分如果你直接上来就写张量级的自动微分大概率会在维度对齐和广播机制上卡住。我的建议是先用纯Python写一个标量级的自动微分引擎大概两百行左右就能跑通。这个阶段的目标不是性能而是让你彻底理解反向传播的链式法则在代码层面是怎么落地的。核心思路是每个标量值包装成一个节点节点记录三个东西数据、梯度、以及产生它的操作和父节点。前向计算时构建计算图反向传播时沿着图反向遍历用链式法则累乘局部梯度。这里有一个容易忽略的细节梯度累加而不是覆盖。因为一个节点可能被多个下游节点使用它的梯度需要把所有路径的贡献加起来。class Value: def __init__(self, data, parents(), op): self.data data self.grad 0.0 self._backward lambda: None self._parents parents self._op op def __add__(self, other): other other if isinstance(other, Value) else Value(other) out Value(self.data other.data, (self, other), ) def _backward(): self.grad out.grad other.grad out.grad out._backward _backward return out上面这段代码就是整个引擎的骨架。你可以在此基础上继续实现乘法、幂运算、激活函数等。写完标量版本之后你会对计算图的拓扑结构和梯度流动有非常直观的认识。这个认知在后续调试张量版本时会反复帮到你。2.2 张量化改造中的广播陷阱从标量过渡到张量最大的坑就是广播机制的反向传播。前向计算时形状为(3, 1)和(1, 4)的两个张量相加会广播成(3, 4)但反向传播时梯度需要沿着被广播的维度求和还原回原始形状。如果你忘了这一步梯度形状就会和参数形状对不上轻则报错重则静默计算出错误结果。我当时的做法是在每个算子的反向函数里先检查输出梯度形状和输入形状是否一致不一致就沿着扩展的维度做sum并keepdims。这个检查逻辑虽然增加了一点开销但在开发阶段帮我省下了大量排查时间。实测下来广播反向的正确处理是从零实现自动微分时最容易出错的地方没有之一。提示在张量自动微分实现中建议为每个算子写一个形状断言函数在开发模式下开启生产模式下关闭。这个习惯能帮你提前拦截百分之八十以上的维度错误。2.3 计算图的构建时机与内存回收计算图有两种构建方式动态图和静态图。从零实现时我建议先做动态图因为调试直观每一步的计算结果都能立刻看到。但动态图的问题是每次前向传播都会重新构建图带来额外的内存分配开销。一个折中方案是使用“磁带”机制前向传播时把所有操作记录到一个线性列表里反向传播时倒序遍历这个列表。这样既保留了动态图的灵活性又避免了递归遍历图结构的栈开销。磁带在每轮迭代结束后需要手动清空否则内存会持续增长。我在早期版本中忘了清空磁带跑了几百轮之后内存直接爆掉这个教训值得你记一下。3. 算子实现矩阵乘法与卷积的手写路径3.1 矩阵乘法的分块优化思路矩阵乘法是AI计算中最核心的算子没有之一。从零实现时最朴素的写法就是三重循环但那个性能在稍大一点的矩阵上就完全不可用。你需要引入分块思想把大矩阵切成小块让每个小块能放进缓存减少内存访问次数。具体来说假设缓存能容纳B x B的子矩阵那么就把A和B分别按B大小分块然后做块间乘法。这个思路在CPU上效果显著在GPU上则对应着共享内存的利用。我实测过一个512 x 512的矩阵乘法朴素三重循环耗时约2.3秒分块之后降到0.4秒左右提升非常明显。def matmul_blocked(A, B, block_size32): n, m len(A), len(B[0]) k len(B) C [[0.0] * m for _ in range(n)] for i0 in range(0, n, block_size): for j0 in range(0, m, block_size): for k0 in range(0, k, block_size): for i in range(i0, min(i0 block_size, n)): for j in range(j0, min(j0 block_size, m)): acc 0.0 for kk in range(k0, min(k0 block_size, k)): acc A[i][kk] * B[kk][j] C[i][j] acc return C这段代码的关键在于最内层循环只访问连续内存缓存命中率高。你可以根据自己机器的缓存大小调整block_size一般取32或64比较稳妥。3.2 卷积算子的im2col展开卷积的实现方式有很多种从零手写时我推荐先用im2col加矩阵乘法的方式。思路是把输入特征图的每个感受野展开成一行形成一个大的矩阵然后和卷积核矩阵做一次矩阵乘法。这样做的好处是复用了你已经优化好的矩阵乘法算子不需要单独为卷积写复杂的循环。im2col的代价是内存占用会增加因为输入数据被重复展开了。一个3 x 3的卷积核展开后内存大约膨胀9倍。对于小批量推理来说可以接受但如果内存紧张就需要考虑原地计算或者分块展开。我在处理较大输入时会把im2col和矩阵乘法都做分块控制峰值内存。反向传播时im2col对应的操作是col2im需要把梯度从展开矩阵散射回原始特征图的位置。这里要注意重叠区域的梯度累加因为同一个输入像素可能被多个感受野覆盖。这个累加逻辑如果写错梯度会偏小训练时表现为收敛变慢但不会报错排查起来相当隐蔽。3.3 激活函数与归一化的数值稳定性ReLU、Sigmoid、Softmax这些激活函数看起来简单但数值稳定性问题非常普遍。Sigmoid在输入绝对值很大时会饱和梯度趋近于零Softmax在指数运算时容易溢出。从零实现时这些细节必须处理。Softmax的标准做法是先减去最大值再取指数这样指数部分的最大值就是0不会溢出。这个技巧在框架里是默认实现的但你自己写的时候如果忘了遇到大数值输入就会得到一堆NaN。我建议在实现每个激活函数时都先问自己一句极端输入下会发生什么归一化层也是类似。批归一化在训练和推理时的行为不同训练时用当前批次的均值和方差推理时用滑动平均。从零实现时滑动平均的更新系数需要仔细调太小则跟不上分布变化太大则波动明显。一般取0.9到0.99之间比较合适。4. 内存管理与推理调度让模型跑得稳4.1 内存池的预分配策略从零搭建推理引擎时如果每次前向传播都动态申请和释放内存性能会被内存分配器拖垮。我的做法是实现一个简单的内存池在初始化阶段预分配一大块连续内存后续所有张量的存储都从这块内存里切分。内存池的关键在于碎片管理。如果张量大小不一反复申请释放会产生外部碎片。一个实用的策略是按大小分桶比如小于1KB的走小桶1KB到1MB的走中桶大于1MB的直接单独分配。每个桶内部用空闲链表管理。这样既能控制碎片又能保证分配速度。注意内存池的预分配大小需要根据你的模型峰值内存来定。太小会频繁触发扩容太大会浪费资源。建议先用动态分配跑一遍记录峰值内存然后按峰值的1.2倍来预分配。4.2 算子调度的拓扑排序当你把模型拆成多个算子之后需要一个调度器来决定算子的执行顺序。最直接的方式是按照计算图的拓扑排序来执行。但拓扑排序只保证依赖关系正确不保证性能最优。一个优化思路是算子融合把逐元素操作如ReLU、加法融合到前一个算子的输出阶段避免中间结果的写回和读取。这个优化在内存带宽受限的场景下效果非常明显。我实测过一个包含二十个逐元素操作的模型融合之后推理延迟降低了约百分之三十五。调度器还需要处理内存复用如果两个张量的生命周期不重叠它们可以共用同一块内存。这需要你分析每个张量的首次使用时间和最后使用时间然后做区间调度。这个逻辑实现起来不复杂但收益很大尤其是对于层数较深的模型。4.3 批处理与动态形状的处理推理服务通常需要支持批处理来提高吞吐但批大小是动态的。从零实现时如果为每个可能的批大小都编译一个计算图维护成本太高。我的做法是支持动态形状算子在执行时根据实际输入形状计算输出形状内存池按最大可能形状预分配。动态形状带来的一个问题是内存池的浪费因为预分配是按最大形状来的。折中方案是设置几个档位比如批大小1、4、8、16各一档请求来时选择最接近的档位。这样既避免了频繁重编译又控制了内存浪费。实测下来档位数量取4到6个比较平衡。5. 从零实现中的调试与验证方法5.1 梯度检查的标准化流程自己实现的自动微分最怕的就是梯度算错了但不知道。梯度检查是必须做的验证步骤。标准做法是用数值微分来近似真实梯度对每个参数施加一个极小的扰动计算损失变化然后和反向传播得到的梯度对比。数值微分的步长选择很关键。太大则近似误差大太小则浮点精度不够。一般取1e-4到1e-5之间。对比时用相对误差而不是绝对误差因为不同参数的梯度量级可能差很多。相对误差小于1e-5算通过小于1e-3算勉强可用再大就说明有问题。我建议把梯度检查写成一个自动化脚本每次修改算子实现后都跑一遍。这个习惯在早期帮我抓到了好几个隐蔽的广播反向错误。5.2 逐层输出对比法当你怀疑某个算子的前向计算有问题时逐层输出对比是最有效的排查手段。具体做法是用你的实现和某个成熟框架分别跑同一个输入然后逐层对比输出。第一个出现显著差异的层就是问题所在。对比时要注意浮点误差的累积。浅层的小误差传到深层可能被放大所以不要要求每一层都完全一致。我的经验是相对误差在1e-6以内算正常超过1e-4就需要检查。如果某一层突然从1e-6跳到1e-3那基本可以锁定问题。5.3 性能剖析的切入点从零实现的代码性能通常不如成熟框架但你需要知道瓶颈在哪。最简单的剖析方法是在每个算子前后打时间戳统计各算子的耗时占比。如果某个算子占了总时间的百分之八十以上那就是优化重点。常见的性能问题包括内存分配过于频繁、循环没有向量化、缓存命中率低。针对性地优化这些点通常能把性能提升一个数量级。但要注意优化之前先确保正确性不要为了性能牺牲正确性。6. 工程化落地中的经验与取舍6.1 什么时候该用现成库从零实现是为了学习但真正上生产时你需要判断哪些部分自己写、哪些部分用现成库。我的原则是核心计算链路自己掌控辅助功能尽量用库。比如矩阵乘法如果你能写到接近BLAS的性能那自己写没问题如果差得远那就调BLAS。但计算图的构建、内存管理、调度逻辑这些和你的业务强相关的部分自己写更可控。6.2 测试覆盖的优先级从零实现的代码测试覆盖要分优先级。最高优先级是梯度正确性因为梯度错了整个训练就是白跑。其次是形状推导形状错了会直接报错反而容易发现。再次是数值稳定性极端输入下的行为需要专门测试。最后是性能回归确保优化没有引入退化。6.3 版本迭代中的兼容性从零实现的过程中接口会不断变化。我的建议是尽早冻结核心接口比如张量的创建、算子的调用签名、内存池的分配释放。这些接口一旦定下来后续的优化和扩展都在内部进行不影响上层调用。这样你的测试用例和示例代码不需要频繁修改。我在早期版本中频繁改动接口导致之前的测试全部失效浪费了不少时间。后来学乖了先把接口定死内部随便改效率高了很多。6.4 文档与注释的取舍从零实现的代码注释要写清楚“为什么”而不是“是什么”。比如广播反向为什么要做sum内存池为什么要分桶这些决策背后的理由比代码本身更重要。过几个月回头看你可能会忘记当时的思路但注释里的理由能帮你快速恢复上下文。文档不需要多但关键设计决策要记录。我习惯在代码仓库里维护一个DESIGN.md记录每个模块的设计动机和权衡。这个习惯在团队协作时尤其重要别人接手你的代码时能快速理解你的意图。7. 从零实现之后的能力迁移当你完整走了一遍从零实现的过程再回头看成熟框架的源码会发现很多东西变得清晰了。你知道计算图是怎么构建的知道梯度是怎么流动的知道内存是怎么管理的。这种认知带来的直接好处是遇到问题时你知道该往哪个方向排查优化时你知道瓶颈可能在哪。更重要的是你获得了一种“拆解”的能力。面对任何新的AI系统或框架你都能快速拆解出它的核心模块理解它的设计取舍。这种能力比会调几个API要值钱得多。我在实际工作中发现从零实现过一遍的工程师在排查线上问题时平均耗时比没实现过的少一半以上。因为他们对底层的直觉更准知道哪些地方容易出问题哪些地方可以放心。最后分享一个小技巧如果你时间有限不需要把每个算子都手写一遍。挑矩阵乘法、卷积、Softmax、批归一化这四个最核心的认真实现并验证。这四个覆盖了大部分计算模式吃透它们其他的算子都是类似的套路。