从零搭建AI工程体系:手写张量、推理服务与性能优化实战

📅 发布时间:2026/10/2 15:48:53
从零搭建AI工程体系:手写张量、推理服务与性能优化实战
1. 从零搭建AI工程体系为什么我劝你别一上来就调包ai-engineering-from-scratch这个标题第一次看到的时候我愣了一下。市面上讲AI的教程铺天盖地但绝大多数都在教你import torch然后跑一个预训练模型真正愿意从工程地基开始讲起的内容少得可怜。我自己带过几个刚入行的同事发现一个很普遍的现象他们能背出Transformer的结构图却说不清楚一个推理服务从请求进来到结果返回中间到底经过了哪些环节、每个环节的瓶颈在哪里、显存是怎么被吃掉的。这就是典型的会用不会造一旦线上出问题就抓瞎。所谓AI工程我的理解是把一个算法原型变成能稳定跑在生产环境里的系统。它跟纯做模型研究是两码事。研究关心的是准确率能不能再涨0.5个点工程关心的是这个模型在QPS 500的时候延迟会不会飙到两秒、显存会不会OOM、服务挂了能不能自动拉起。从零搭建这套东西核心价值不在于让你重新发明轮子而在于让你真正理解每个轮子为什么长这样。当你亲手写过一个最朴素的推理循环再去看那些框架的源码很多设计决策就豁然开朗了。这篇内容适合谁看如果你是有一定Python基础、想往AI工程方向转的开发者或者已经在做算法但想补齐工程能力的同学再或者你是团队里负责把模型落地的那个人那接下来的内容应该对你有用。我会从最底层的张量操作讲起一路搭到推理服务、性能优化和监控全程不依赖任何高级封装能自己写的就自己写。当然该用现成库的地方我也会说清楚为什么用、用哪个、怎么选。需要提前说明的是从零搭建不等于什么都手写。工程的核心是权衡是知道什么时候该造轮子、什么时候该用轮子。我会在每个环节解释清楚这个权衡的逻辑这样你以后遇到新场景也能自己判断。2. 整体架构设计与技术选型思路2.1 为什么选择自底向上的搭建路径我见过太多人学AI工程的方式是反过来的先学怎么用FastAPI包一个接口再学怎么用ONNX导出模型最后才回头补张量和算子的知识。这种路径的问题是你始终在知其然的层面打转遇到框架没覆盖的场景就束手无策。自底向上的路径虽然前期慢但后劲足。具体来说我设计的搭建顺序是这样的先搞明白数据在内存里怎么表示张量基础再搞明白计算怎么发生前向传播的手写实现然后是计算怎么加速向量化、批处理接着是模型怎么存和怎么读序列化最后才是服务怎么对外提供API层、并发、监控。这个顺序的每一层都建立在前一层的基础上不会出现知识断层。提示如果你时间有限可以跳过手写前向传播的部分直接看服务化但我强烈建议至少把张量内存布局那一节看完那是后面所有性能优化的基础。2.2 技术栈的取舍哪些自己写哪些用现成的从零搭建最容易走偏的地方就是什么都想自己写最后写出来一个又慢又难维护的四不像。我的原则是核心链路上自己写边缘设施用成熟方案。核心链路指的是张量运算、推理调度、批处理逻辑这些直接决定性能和正确性的部分。这些自己写是因为你需要对每一行代码的开销心里有数。边缘设施指的是日志、监控、配置管理、进程守护这些跟AI本身关系不大的部分用现成的库能省下大量时间而且它们足够成熟没必要重复造。模块自己实现用现成方案选择理由张量运算是-理解内存布局和计算开销推理调度是-控制批处理和并发行为模型序列化部分safetensors格式标准自己写解析器性价比低API服务是-逻辑简单自己写更可控日志监控否logging prometheus成熟稳定没必要重写配置管理否pydantic-settings类型校验省心这个表格背后的逻辑其实很简单凡是影响模型跑得快不快、对不对的自己写凡是影响运维方不方便的用现成的。因为前者是你的核心竞争力所在后者是通用需求社区已经打磨得很好了。2.3 目录结构规划让工程有工程的样子从零搭建的另一个好处是你可以从一开始就把目录结构设计好而不是等项目膨胀了再重构。我习惯的结构是这样的ai-eng-scratch/ ├── core/ # 核心张量与算子 │ ├── tensor.py │ └── ops.py ├── model/ # 模型定义与加载 │ ├── layers.py │ └── loader.py ├── serving/ # 推理服务 │ ├── engine.py │ ├── batching.py │ └── api.py ├── utils/ # 工具函数 │ ├── logging.py │ └── metrics.py ├── configs/ # 配置文件 └── tests/ # 测试这个结构的关键在于core和serving的分离。core是纯计算不依赖任何网络或IOserving负责把core包装成服务。这样分层的好处是core可以单独测试也可以被其他程序复用不会因为服务层的改动而受影响。很多项目把计算逻辑和服务逻辑混在一起最后想换个服务框架就得大动干戈这就是没分层的代价。3. 核心细节解析与实操要点3.1 张量内存布局一切性能问题的根源很多人觉得张量就是个多维数组用numpy或者list套list不就行了。这个理解在功能层面没错但在性能层面差得远。张量的核心在于连续内存 步长stride这套机制。我举个具体的例子。一个形状为(2, 3)的二维张量在内存里其实是一段连续的12个浮点数假设float32按行优先排列。访问[i][j]的时候实际的内存偏移是i * 3 j。这里的3就是行方向的步长。如果是转置操作传统做法是重新拷贝一份数据但聪明的做法是只改步长把行步长和列步长对调数据一个字节都不动。这就是所谓的视图操作零拷贝。class Tensor: def __init__(self, data, shape, stridesNone): self.data data # 底层连续内存用array或list self.shape shape if strides is None: # 按行优先计算步长 strides [] acc 1 for dim in reversed(shape): strides.insert(0, acc) acc * dim self.strides strides else: self.strides strides def transpose(self): # 只交换步长不拷贝数据 new_shape list(reversed(self.shape)) new_strides list(reversed(self.strides)) return Tensor(self.data, new_shape, new_strides)这段代码看起来简单但它解释了一个关键问题为什么有些操作快有些操作慢。转置是O(1)的因为它只改了元数据而如果要对转置后的张量做逐元素运算就会因为内存访问不连续而变慢这就是所谓的cache不友好。理解了这一点你就能明白为什么深度学习框架里到处都在强调contiguous。注意自己实现张量的时候一定要区分视图和拷贝。视图共享底层内存修改一个会影响另一个这是很多bug的来源。我建议在视图操作上加明确的命名比如transpose_view提醒调用者。3.2 前向传播的手写实现把黑盒拆开看搞清楚了张量下一步就是计算。我拿一个最简单的全连接层加激活函数来演示。别看它简单理解了它的计算过程Transformer里的注意力机制本质上也是同样的东西只是多了几个矩阵乘和softmax。全连接层的计算是y xW b其中x是输入W是权重b是偏置。手写实现的时候最容易忽略的是批处理维度。假设输入x的形状是(batch, in_features)W的形状是(in_features, out_features)那么输出y的形状是(batch, out_features)。矩阵乘法的三重循环写法是这样的def matmul(a, b): # a: (M, K), b: (K, N) - (M, N) M, K a.shape K2, N b.shape assert K K2 result [[0.0] * N for _ in range(M)] for i in range(M): for k in range(K): a_ik a[i][k] if a_ik 0: continue # 稀疏优化 for j in range(N): result[i][j] a_ik * b[k][j] return result这个三重循环是理解矩阵乘法的起点但它的性能很差因为内层循环访问b[k][j]的时候j在变而k不变内存访问是跳跃的。优化的第一步是调整循环顺序把j放到最内层让result[i][j]和b[k][j]都能顺序访问。再进一步就是分块tiling把大矩阵切成能放进CPU缓存的小块这是所有高性能矩阵库的核心技巧。我实测过一个朴素的Python三重循环做1024x1024的矩阵乘法大概要几十秒调整循环顺序后能快两三倍用numpy的话是毫秒级。这个差距就是工程优化的空间所在。你不需要自己写出numpy级别的性能但你需要知道差距在哪里、怎么缩小。3.3 批处理与动态形状推理服务的核心难题单个请求的推理很简单难的是同时处理多个请求。批处理的核心思想是把多个请求的输入拼成一个更大的张量一次算完再拆开。这样做的好处是充分利用硬件的并行能力坏处是不同请求的输入长度可能不一样拼不起来。解决这个问题有几种常见策略。第一种是padding把所有输入补齐到最大长度缺点是浪费计算。第二种是分桶bucketing把长度相近的请求分到一组减少padding浪费。第三种是动态批处理dynamic batching服务端攒一小段时间的请求再一起处理用时间换吞吐。class DynamicBatcher: def __init__(self, max_batch_size32, max_wait_ms10): self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue [] def add_request(self, request): self.queue.append(request) if len(self.queue) self.max_batch_size: return self.flush() return None def flush(self): if not self.queue: return None batch self.queue[:self.max_batch_size] self.queue self.queue[self.max_batch_size:] return batch这里的max_wait_ms是个关键参数。设太小攒不够请求吞吐上不去设太大用户等得久延迟变高。我一般从10ms开始调根据实际压测结果微调。这个参数没有标准答案取决于你的业务对延迟的敏感程度。提示动态批处理在GPU上收益明显在CPU上收益有限因为CPU本身并行能力就弱。如果你的服务跑在CPU上优先考虑用多进程而不是批处理。4. 实操过程与核心环节实现4.1 环境准备与依赖安装动手之前先把环境弄干净。我强烈建议用虚拟环境别在系统Python里瞎装。conda或者venv都行我个人偏好venv轻量。python -m venv ai-eng source ai-eng/bin/activate # Windows用 ai-eng\Scripts\activate pip install numpy safetensors fastapi uvicorn pydantic-settings prometheus-client这里解释一下每个依赖的用途。numpy不用多说做数值计算的基础safetensors是模型权重序列化的格式比pickle安全加载也快fastapi和uvicorn搭API服务pydantic-settings管配置prometheus-client做指标暴露。注意我没有装torch或tensorflow因为从零搭建的核心就是不依赖这些重型框架。装完之后验证一下python -c import numpy; print(numpy.__version__)能打印出版本号就说明环境没问题。如果报错大概率是虚拟环境没激活检查一下命令行前面有没有(ai-eng)的提示。4.2 手写张量类的完整实现前面讲了张量的原理这里给一个能跑的最小实现。我把它放在core/tensor.py里。import numpy as np class Tensor: def __init__(self, data, shapeNone, stridesNone): if isinstance(data, np.ndarray): self.data data self.shape shape or data.shape else: self.data np.array(data, dtypenp.float32) self.shape shape or self.data.shape self.strides strides or self._compute_strides(self.shape) staticmethod def _compute_strides(shape): strides [] acc 1 for dim in reversed(shape): strides.insert(0, acc) acc * dim return strides def reshape(self, new_shape): total 1 for d in new_shape: total * d assert total self.size(), reshape前后元素数量必须一致 return Tensor(self.data.reshape(new_shape)) def size(self): result 1 for d in self.shape: result * d return result def __matmul__(self, other): return Tensor(self.data other.data) def __add__(self, other): return Tensor(self.data other.data) def relu(self): return Tensor(np.maximum(self.data, 0))这个实现基于numpy但封装了形状和步长的管理。你可能会问既然用了numpy为什么不直接用numpy的ndarray因为我要的是一个统一的接口后面加自定义算子、加设备管理CPU/GPU的时候这层封装就是扩展点。直接裸用numpy后面想换后端就得改遍所有代码。实测下来这个封装带来的性能损耗可以忽略不计因为真正的计算还是numpy在做。封装层只做参数校验和形状推导开销很小。4.3 模型加载与权重管理模型权重通常存成键值对的形式键是层名值是权重张量。safetensors的格式很简洁加载起来也快。from safetensors.numpy import load_file class ModelLoader: def __init__(self, path): self.weights load_file(path) def get_layer(self, name): if name not in self.weights: raise KeyError(f权重 {name} 不存在可用的有{list(self.weights.keys())}) return Tensor(self.weights[name]) def summary(self): total_params 0 for name, w in self.weights.items(): print(f{name}: shape{w.shape}, dtype{w.dtype}) total_params w.size print(f总参数量: {total_params})加载的时候有个坑要注意safetensors默认返回的是只读数组如果你想在加载后做量化或者微调得先拷贝一份。我踩过这个坑当时想对权重做归一化结果报了个只读错误排查了半天。注意加载大模型的时候内存峰值可能是模型大小的两倍因为加载过程中旧数据和新数据会短暂共存。如果内存紧张考虑分片加载。4.4 推理引擎的搭建推理引擎是核心中的核心它负责把输入数据喂给模型、执行计算、返回结果。我设计成一个类把批处理、模型调用、结果后处理都串起来。class InferenceEngine: def __init__(self, model, batcher): self.model model self.batcher batcher def predict(self, inputs): # inputs: list of Tensor batch self.batcher.add_request(inputs) if batch is None: return None # 还没攒够等下一轮 return self._run_batch(batch) def _run_batch(self, batch): # 把多个输入拼成一个批次张量 stacked np.stack([x.data for x in batch], axis0) batch_tensor Tensor(stacked) # 前向传播 output self.model.forward(batch_tensor) # 拆回单个结果 return [Tensor(output.data[i]) for i in range(len(batch))]这里的np.stack要求所有输入形状一致如果不一致就得先padding。padding的逻辑我单独抽成一个函数因为不同任务的padding策略不一样文本任务补0图像任务可能补边缘像素。4.5 API服务与并发处理最后一步是把引擎包装成HTTP服务。用FastAPI写起来很简洁但要注意并发模型。FastAPI默认是异步的而我们的推理是同步的CPU密集操作直接放在async函数里会阻塞事件循环。from fastapi import FastAPI from concurrent.futures import ThreadPoolExecutor import asyncio app FastAPI() executor ThreadPoolExecutor(max_workers4) engine None # 初始化时赋值 app.post(/predict) async def predict(payload: dict): inputs [Tensor(payload[data])] loop asyncio.get_event_loop() result await loop.run_in_executor(executor, engine.predict, inputs) if result is None: return {status: queued} return {result: result[0].data.tolist()}max_workers设多少合适经验值是CPU核心数或者核心数的1.5倍。设太多会导致上下文切换开销设太少又浪费CPU。我一般先用os.cpu_count()然后压测调整。5. 常见问题与排查技巧实录5.1 性能问题的排查思路推理服务慢原因可能有很多层。我的排查顺序是从外到内先看网络和框架层再看批处理逻辑最后看计算本身。现象可能原因排查方法解决方向延迟高但CPU不高批处理等待时间过长看max_wait_ms配置调小等待时间CPU跑满但吞吐低计算未向量化profile热点函数用numpy替换循环内存持续增长张量未释放看对象引用及时del和gc偶发超时大请求阻塞小请求看请求大小分布分队列处理这个表是我踩坑总结出来的基本覆盖了八成的情况。特别说一下内存增长那个Python的垃圾回收对循环引用处理得不好如果你在张量之间建立了循环引用内存就回收不掉。我一般会在服务里定期手动gc.collect()虽然有点土但管用。5.2 数值精度问题从零实现的时候精度问题特别容易出。numpy默认是float64而模型权重通常是float32混用会导致结果对不上。我建议全程统一用float32在创建张量的时候就指定dtype。# 错误示范 x Tensor([1.0, 2.0, 3.0]) # 默认float64 # 正确做法 x Tensor(np.array([1.0, 2.0, 3.0], dtypenp.float32))还有一个隐蔽的坑是累加顺序。浮点加法不满足结合律(ab)c和a(bc)的结果可能差一点点。在矩阵乘法里不同的循环顺序会导致不同的累加顺序结果就有微小差异。这个差异在单次计算里可以忽略但在需要严格复现的场景比如对比两个实现就会成为问题。我的做法是固定循环顺序并且在测试里用np.allclose而不是来比较。5.3 并发安全的坑推理引擎如果被多个线程同时调用共享状态就会出问题。最常见的是批处理队列多个线程同时往队列里加请求不加锁就会丢数据或者重复处理。import threading class SafeBatcher: def __init__(self, max_batch_size32): self.max_batch_size max_batch_size self.queue [] self.lock threading.Lock() def add_request(self, request): with self.lock: self.queue.append(request) if len(self.queue) self.max_batch_size: batch self.queue[:] self.queue [] return batch return None锁的粒度要控制好只锁队列操作不要把整个推理过程都锁住否则并发就退化成串行了。我见过有人把锁加在predict方法上结果QPS上不去排查半天才发现是锁的问题。提示Python的GIL让很多人以为不用加锁但GIL只保证字节码级别的原子性像list.append这种操作在GIL切换的间隙仍可能出问题。涉及共享可变状态该加锁就加锁。5.4 模型版本管理服务上线后模型会迭代怎么在不重启服务的情况下切换模型是个实际问题。我的做法是维护一个模型注册表新模型加载好之后原子替换引用。class ModelRegistry: def __init__(self): self._models {} self._active None self.lock threading.Lock() def register(self, version, model): with self.lock: self._models[version] model def activate(self, version): with self.lock: if version not in self._models: raise KeyError(f版本 {version} 未注册) self._active self._models[version] def get_active(self): return self._active切换的时候要注意正在处理的请求可能还在用旧模型不能立刻把旧模型释放掉。我一般会等一个宽限期确认没有请求在用旧模型了再清理。6. 从能跑到好用还差哪些工程细节6.1 监控指标的埋点服务上线不是终点能观测才是。我至少会埋这几个指标请求总数、请求延迟分布、批处理大小分布、错误率。用prometheus-client暴露出来配合Grafana看板。from prometheus_client import Counter, Histogram REQUEST_COUNT Counter(inference_requests_total, 总请求数) REQUEST_LATENCY Histogram(inference_latency_seconds, 请求延迟) BATCH_SIZE Histogram(inference_batch_size, 批处理大小, buckets[1, 2, 4, 8, 16, 32, 64]) def predict_with_metrics(engine, inputs): REQUEST_COUNT.inc() with REQUEST_LATENCY.time(): result engine.predict(inputs) return result延迟的直方图分桶要合理我一般按2的幂次分因为延迟的分布通常是对数正态的。分桶太粗看不出细节太细又浪费存储。6.2 优雅关闭与请求排空服务重启的时候如果直接kill正在处理的请求就丢了。正确的做法是收到关闭信号后停止接受新请求等现有请求处理完再退出。import signal import sys shutdown_flag False def handle_signal(signum, frame): global shutdown_flag shutdown_flag True print(收到关闭信号等待请求排空...) signal.signal(signal.SIGTERM, handle_signal) signal.signal(signal.SIGINT, handle_signal)然后在请求处理循环里检查shutdown_flag为True就不再接新请求。这个逻辑看起来简单但很多服务都忽略了导致每次发布都有一批用户请求失败。6.3 配置的热更新有些配置改了不该重启服务比如批处理大小、日志级别。我用pydantic-settings配合文件监听实现热更新。from pydantic_settings import BaseSettings import json class Settings(BaseSettings): max_batch_size: int 32 max_wait_ms: int 10 log_level: str INFO def reload_settings(pathconfigs/settings.json): with open(path) as f: data json.load(f) return Settings(**data)热更新的关键是原子性新配置解析失败要保留旧配置不能把服务搞挂。我一般会先校验新配置校验通过再替换。6.4 压测与容量规划上线前一定要压测知道服务的极限在哪里。我用locust或者wrk模拟不同并发下的表现。wrk -t4 -c100 -d30s --latency http://localhost:8000/predict-t4是4个线程-c100是100个并发连接-d30s跑30秒。看结果里的P99延迟和QPSP99超过业务容忍度就得优化。容量规划的时候留30%的余量别把服务跑到极限否则流量一波动就雪崩。我个人的经验是从零搭建这套东西最大的收获不是代码本身而是对每个环节的手感。你知道一个矩阵乘法大概要多久知道批处理能带来多少吞吐提升知道内存什么时候会爆。这种手感是调包调不出来的只能自己动手才能建立。后面你再去看那些框架的文档会发现很多之前看不懂的设计说明现在一眼就明白了。这个内容后续还可以往分布式推理、模型量化、异构硬件调度这些方向扩展但那是另一个话题了。