不调包手搓AI工程:从零实现分词、注意力与推理服务

发布时间:2026/11/3 23:13:09
不调包手搓AI工程:从零实现分词、注意力与推理服务 从零手搓AI工程一个“不调包”主义者的项目复盘这两年AI应用层的工具链成熟得有点吓人。LangChain、LlamaIndex、各种Agent框架轮番上阵好像不 import 几个库都不好意思说自己做过AI项目。但我越来越觉得真正让人心里发虚的不是“不会用框架”而是把框架扒掉之后脑子里一片空白——token怎么切、向量怎么算、注意力到底在算什么、推理服务怎么扛住并发这些最底层的东西全被封装成了黑盒。ai-engineering-from-scratch这个项目就是冲着这个焦虑去的。它的核心主张很朴素不依赖高层框架用最基础的Python和数学把AI工程里那些关键组件一个一个手写出来。适合谁看适合已经会用现成工具、但想搞清楚“里面到底发生了什么”的开发者也适合刚入门、不想一上来就被框架绑架的新人。这篇复盘我会把项目里几个核心模块的拆解思路、踩过的坑、以及那些文档里不会写的经验尽量讲透。1. 为什么我坚持先手写一遍再谈调包1.1 框架带来的“能力幻觉”我先说个真实感受。早些年我用某个流行框架搭RAG三行代码就把文档灌进去、问题答出来了当时觉得自己挺牛。直到有一次检索结果明显跑偏我盯着那几行配置愣了半天完全不知道问题出在切分粒度、embedding模型还是相似度计算上。因为整条链路都被封装了我连中间张量的形状都看不到。这就是能力幻觉你以为你掌握了其实你只是会调用。框架本身没问题它极大提升了效率但如果你从没手写过一遍你就不具备“定位问题”的能力。ai-engineering-from-scratch的出发点就是打破这个幻觉——每个组件都自己实现哪怕实现得粗糙至少你知道每一行在干什么。1.2 手写一遍到底能获得什么具体来说手写能带来三个层面的收益这也是我在项目里反复验证的可调试性当检索不准时你能打印出query向量和doc向量的余弦相似度分布一眼看出是归一化没做还是维度对不上。可替换性你知道每个环节的输入输出契约换模型、换算法时心里有底不会因为框架版本升级就崩。性能直觉手写过一遍你才知道一次矩阵乘法有多贵、为什么batch size会影响显存、为什么KV Cache能省那么多算力。提示手写不等于生产环境要手写。生产该用框架还是用框架手写的价值在于建立心智模型让你在框架出问题时能往下钻。1.3 项目的整体模块划分ai-engineering-from-scratch大致分成这么几块我按依赖顺序排模块核心内容依赖分词与词表BPE、WordPiece原理与实现纯Python向量化TF-IDF、embedding基础NumPy注意力机制Scaled Dot-Product、多头注意力NumPy迷你Transformer前向传播、简单训练循环NumPy检索系统向量索引、相似度检索NumPy推理服务批处理、缓存、并发标准库轻量Web这个划分的好处是层层递进分词是文本进模型的入口向量化是检索的基础注意力是Transformer的心脏检索和服务则是把前面所有东西串成一个能用的系统。每一块都能独立跑也能组合起来。2. 分词器所有文本任务的第一道关2.1 为什么分词比想象中重要很多人觉得分词就是个预处理随便调个库就行。但实际项目里分词直接决定了模型的词表大小、序列长度、以及OOV未登录词的处理方式。我见过一个中文项目因为用了按字切分的方案序列长度直接翻倍推理成本涨了一大截也见过英文项目因为BPE的merge规则没调好把常见词切得稀碎。ai-engineering-from-scratch里我实现了一个简化版BPE核心逻辑就三步统计相邻字符对频率、合并最高频的pair、重复直到达到目标词表大小。听起来简单但细节全是坑。2.2 手写BPE的关键步骤先看核心循环的伪代码思路def train_bpe(corpus, vocab_size): # 初始化把每个词拆成字符序列末尾加结束符 vocab build_initial_vocab(corpus) merges [] while len(vocab) vocab_size: # 1. 统计所有相邻pair的频率 pair_freq count_pairs(corpus) if not pair_freq: break # 2. 选最高频的pair best_pair max(pair_freq, keypair_freq.get) # 3. 合并 corpus merge_pair(corpus, best_pair) merges.append(best_pair) vocab.add(.join(best_pair)) return vocab, merges这里有几个必须注意的点都是我自己踩出来的结束符不能省每个词末尾加一个特殊符号比如/w否则low和lower的切分会混乱因为low可能是lower的前缀。频率统计要基于当前语料状态每合并一次语料就变了必须重新统计不能缓存。合并顺序要记录推理时要用同样的merge顺序否则训练和推理不一致效果直接崩。2.3 训练与推理的一致性陷阱这是最容易翻车的地方。训练时你得到了一串merge规则推理时必须按完全相同的顺序应用。我一开始图省事推理时直接按词表做最长匹配结果同一个词训练时切成[lo, w]推理时切成[low]embedding对不上模型表现一塌糊涂。正确的做法是把merges存成一个有序列表推理时从头遍历能合并就合并。虽然慢一点但保证了一致性。实测下来一个几万词的词表单条文本分词耗时在毫秒级完全可接受。注意如果你的语料里有大量数字或特殊符号建议先做归一化比如全角转半角、统一大小写否则词表会被这些噪声撑爆。3. 向量化与相似度检索系统的地基3.1 从TF-IDF到稠密向量检索系统的核心是“把文本变成能比较的向量”。最经典的是TF-IDF简单、可解释、不需要训练。ai-engineering-from-scratch里我先实现了TF-IDF再过渡到稠密向量目的是让读者理解两者的本质区别。TF-IDF的逻辑是一个词在当前文档出现越多TF高在整个语料出现越少IDF高它就越有代表性。实现上就是统计词频、算逆文档频率、然后加权。但它的硬伤是词汇鸿沟——电脑和计算机在TF-IDF眼里是完全无关的两个维度语义相似度为零。稠密向量embedding解决了这个问题它把文本映射到一个低维连续空间语义相近的文本距离近。但embedding需要模型而ai-engineering-from-scratch的原则是尽量不依赖预训练模型所以我用了一个基于共现矩阵的简化embedding统计词与词在窗口内的共现次数做SVD降维得到每个词的向量。虽然效果比不上真正的神经embedding但足以讲清楚原理。3.2 余弦相似度的实现细节有了向量下一步是算相似度。最常用的是余弦相似度def cosine_similarity(a, b): dot np.dot(a, b) norm_a np.linalg.norm(a) norm_b np.linalg.norm(b) if norm_a 0 or norm_b 0: return 0.0 return dot / (norm_a * norm_b)看起来没毛病但实际用的时候有两个坑零向量问题如果某个文档全是停用词向量可能全零除零会报错。必须加保护。归一化预计算如果要对一个query和上万个doc算相似度每次都算norm_b是浪费。正确做法是提前把所有doc向量归一化检索时只算点积速度能快好几倍。我实测过1万个768维向量不预归一化算一次全量检索要几百毫秒预归一化后降到几十毫秒。这个优化在数据量上来之后是质变。3.3 向量索引的朴素实现与优化数据量小的时候暴力遍历所有向量算相似度完全够用。但数据量上万之后就得考虑索引。ai-engineering-from-scratch里我实现了一个基于聚类的简化索引类似IVF的思路先用K-Means把所有doc向量聚成N个簇每个簇有个中心向量。检索时先算query和所有簇中心的相似度选出最近的几个簇。只在这几个簇内部做暴力检索。这样能把检索范围缩小到原来的几分之一。代价是可能漏掉一些边界情况query正好在两个簇之间所以簇的数量和每簇检索数需要调。我的经验是簇数量取sqrt(总向量数)左右每簇检索数取总簇数的10%到20%能在召回率和速度之间取得不错的平衡。方案召回率检索速度适用规模暴力遍历100%慢1万聚类索引90%-95%快1万-100万图索引(HNSW)95%-99%很快百万级以上4. 注意力机制亲手推导一遍才叫懂4.1 从“加权平均”理解注意力注意力机制被讲烂了但很多人还是似懂非懂。我的理解方式很简单注意力就是一种带权重的平均。你有一堆value每个value有个权重权重由query和key的相似度决定。相似度高的value权重就大最后加权求和。用生活类比你去图书馆找书query每本书有个标签key书的内容是value。你拿自己的需求去和每个标签比对越匹配的书你越重视它的内容。最后你脑子里形成的印象就是所有书内容的加权平均。ai-engineering-from-scratch里我手写了Scaled Dot-Product Attentiondef attention(Q, K, V): d_k Q.shape[-1] # 1. Q和K做点积得到相似度分数 scores np.matmul(Q, K.transpose(-2, -1)) # 2. 缩放防止softmax梯度消失 scores scores / np.sqrt(d_k) # 3. softmax归一化成权重 weights softmax(scores, axis-1) # 4. 加权求和 return np.matmul(weights, V)4.2 为什么要除以根号d_k这是面试常问、也是实际实现必须理解的细节。假设Q和K的每个分量都是均值0、方差1的独立随机变量那么它们的点积d_k个分量相乘再求和的方差就是d_k。当d_k很大时比如512点积的数值会很大softmax之后会变得极其尖锐——几乎所有的权重都集中在一个位置上梯度接近零训练不动。除以sqrt(d_k)就是把方差拉回到1让softmax的输出分布更平滑梯度更健康。这个操作看起来只是个小trick但它直接决定了深层网络能不能训起来。4.3 多头注意力的拆分与拼接单头注意力只能捕捉一种关系模式多头则是让模型同时关注不同的子空间。实现上就是把Q、K、V分别投影到h个低维空间各自算注意力最后拼接再投影回来。def multi_head_attention(X, W_q, W_k, W_v, W_o, num_heads): # 投影 Q X W_q K X W_k V X W_v # 拆成多头 Q split_heads(Q, num_heads) K split_heads(K, num_heads) V split_heads(V, num_heads) # 各自算注意力 attn attention(Q, K, V) # 拼接 concat merge_heads(attn) # 输出投影 return concat W_o这里最容易错的是维度变换。split_heads的时候要把最后一维拆成(num_heads, head_dim)然后调整轴顺序让head维度在batch之后。我一开始轴顺序搞反了结果每个头拿到的数据是错的训练loss死活不降。排查了半天才发现是transpose的轴写错了。提示手写多头注意力时建议先用小维度比如d_model8, num_heads2跑一遍打印每一步的shape确认无误再放大。shape对不上是这类实现最常见的bug。5. 迷你Transformer把零件组装起来5.1 前向传播的完整链路有了分词、embedding、注意力就可以组装一个迷你Transformer了。ai-engineering-from-scratch里的版本只做前向传播和最简单的训练目的是让读者看清数据怎么流动。完整链路是token ids → embedding查表 → 加位置编码 → N层Transformer Block每层含多头注意力前馈网络残差LayerNorm→ 输出投影 → softmax得到下一个token的概率。位置编码我用了最经典的正弦余弦方案def positional_encoding(seq_len, d_model): pos np.arange(seq_len)[:, None] i np.arange(d_model)[None, :] angle pos / np.power(10000, (2 * (i // 2)) / d_model) pe np.zeros((seq_len, d_model)) pe[:, 0::2] np.sin(angle[:, 0::2]) pe[:, 1::2] np.cos(angle[:, 1::2]) return pe这个方案的好处是不需要训练而且能外推到比训练时更长的序列虽然外推效果有限。5.2 残差连接和LayerNorm为什么不能少残差连接x sublayer(x)解决的是深层网络的梯度消失问题。没有残差梯度在反向传播时经过多层会指数级衰减浅层几乎学不到东西。LayerNorm则是稳定每层的输入分布让训练更平滑。我做过对比实验去掉残差3层网络loss就降不下去了去掉LayerNorm训练初期loss震荡严重。这两个组件看起来简单但它们是深层网络能训起来的基石。5.3 一个能跑通的最小训练循环训练循环的核心就是前向算loss、反向算梯度、更新参数。ai-engineering-from-scratch里我用数值梯度做了个极简版反向传播虽然慢但绝对正确目的是让读者理解梯度的本质。def train_step(model, x, y, lr0.01): # 前向 logits model.forward(x) loss cross_entropy(logits, y) # 数值梯度仅用于教学生产用自动微分 grads numerical_gradients(model, x, y) # 更新 for param, grad in zip(model.params, grads): param - lr * grad return loss数值梯度的思路是对每个参数稍微加一点看loss变化多少变化率就是梯度。慢是真慢但它能帮你验证自动微分的实现对不对。我建议每个手写反向传播的人都用数值梯度校验一遍能省下大量debug时间。6. 推理服务从能跑到能扛6.1 批处理为什么能提升吞吐单条推理时GPU或CPU的算力大量浪费在等待数据搬运上。批处理把多条请求攒在一起算矩阵乘法的规模变大算力利用率显著提升。实测下来batch size从1提到8吞吐能提升3到5倍延迟只增加一点点。但批处理有个矛盾攒批会增加延迟。如果请求来得稀疏等攒够一批可能已经超时了。所以生产系统通常用动态批处理设置一个最大等待时间比如10ms时间到了或者批满了就发出去。6.2 KV Cache省下的算力自回归生成时每生成一个新token都要重新算一遍前面所有token的注意力。但前面token的K和V其实没变重复计算纯属浪费。KV Cache就是把已经算过的K和V存下来生成新token时只算新token的Q和K、V然后和缓存的拼接。这个优化能把生成第n个token的复杂度从O(n²)降到O(n)。序列越长收益越大。我实测过生成512个token用KV Cache比不用快将近10倍。注意KV Cache会占显存序列越长、batch越大占用越多。生产环境要算好显存预算必要时用分页或量化来压缩。6.3 并发下的资源竞争与限流推理服务上线后最怕的是突发流量把资源打满。ai-engineering-from-scratch里我用标准库的队列和线程池做了个简化版服务核心是请求队列固定worker数。请求进来先入队worker从队列取任务处理队列满了就拒绝返回限流响应。这样做的目的是保护后端不被压垮。限流阈值怎么定我的经验是先压测出单worker的QPS乘以worker数再打个7折作为安全阈值。留30%余量应对突发。策略优点缺点直接处理实现简单高并发下崩溃队列限流保护后端需要调参动态扩缩容弹性好复杂度高7. 那些文档里不会写的踩坑经验7.1 浮点精度问题比你想象的常见手写数值计算时浮点误差会累积。我遇到过softmax因为输入数值太大导致溢出exp(1000)直接inf解决办法是减去最大值softmax(x) softmax(x - max(x))数学上等价但数值稳定。还有余弦相似度两个几乎相同的向量算出来可能是0.9999999而不是1.0做阈值判断时要留epsilon。7.2 维度顺序是万恶之源NumPy和PyTorch的默认维度顺序不一样手写时如果混用transpose和reshape会把你搞疯。我的建议是全程用一套约定比如统一用(batch, seq, dim)写注释标清楚每个张量的shape每步都打印验证。这个习惯能省下无数debug时间。7.3 别急着优化先跑通我一开始写检索想着直接上HNSW结果调了两天没调通连基本功能都没验证。后来退回去先写暴力遍历跑通了再逐步加索引反而快得多。先正确再快这个顺序不能反。手写项目的价值在于理解不在于性能别本末倒置。7.4 测试用例要覆盖边界分词要测空字符串、超长文本、纯符号相似度要测零向量、相同向量、正交向量注意力要测单头、多头、不同序列长度。这些边界情况在真实数据里一定会遇到提前测好能省很多事。8. 这个项目适合怎么用如果你是想入门AI工程的开发者我建议按模块顺序走一遍每个模块都自己敲一遍代码别复制粘贴。敲的过程中遇到shape对不上、loss不降正是学习发生的地方。如果你已经有一定经验可以挑自己最薄弱的模块深入比如一直没搞懂注意力就老老实实手推一遍公式再实现。也可以把这个项目当成面试准备材料手写过一遍的人回答底层问题时底气完全不一样。如果你想把它用到实际项目里我的建议是用它的思路不用它的代码。生产环境该用成熟框架就用但心里要清楚每个环节在干什么出问题时知道往哪查。这才是ai-engineering-from-scratch真正的价值——不是替代框架而是让你在框架之上多一层掌控感。最后分享一个我自己的习惯每次用新框架我都会花半天时间用最朴素的方式把它的核心功能手写一遍。这个过程很笨但每次都能让我对这个框架的理解深一层。手写不是为了炫技是为了在抽象泄漏的时候你还能接得住。

关于本文作者

来自尧图内容编辑团队

尧图内容编辑团队 内容团队

尧图内容编辑团队

本文由尧图网络内容编辑团队执笔。团队由资深项目经理、前端工程师与设计师组成,所有内容均来自亲手交付的真实项目,先讲清问题、再给出可落地的解法。尧图深耕北京网站建设十年,服务过京华建材集团、智造科技等各行业客户,把一线经验沉淀为可复用的行业观察。

  • 十年建站经验,覆盖建材、制造、服务、文创等
  • 项目经理把关选题与事实准确性
  • 工程师与设计师联合撰写专业细节
  • 统一编辑规范,保证文风与排版一致
  • 每月复盘转化数据,迭代选题方向

延伸阅读

相关资讯与近期热门内容

深度阅读推荐

建站决策前值得细读的三篇

网站改版的5个关键决策
2024-08-12

网站改版的5个关键决策

什么时候该改版、改到什么程度、如何避免流量掉光,京华建材集团改版复盘给出答案。

获取专属建站方案

看完文章,把您的行业与预算告诉我们,免费获取一份量身定制的官网建设方案与报价。

立即免费咨询