ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

从零手搓AI工程:告别调包,掌握底层训练与推理部署

从零手搓AI工程:告别调包,掌握底层训练与推理部署 1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经掌握了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出和请求超时我才意识到——只会调包的人根本不知道模型在底层到底经历了什么。ai-engineering-from-scratch这个项目标题核心讲的其实就是一件事把AI工程当成一门手艺来学从最底层的张量运算、数据管道、模型训练循环、推理服务部署一层一层自己搭起来。它不是教你如何用某个框架的某个API而是让你理解每一行代码背后的计算逻辑和工程取舍。适合谁看适合那些已经会写Python、用过PyTorch或TensorFlow跑过几个Demo但一遇到性能瓶颈、显存爆炸、数据加载慢、服务不稳定就束手无策的开发者。也适合想转行做AI工程、但被各种“三天速成”课程带偏了方向的人。我自己走过最大的弯路就是早期太依赖高层封装。用Trainer类跑微调用pipeline做推理看起来很爽但一旦要改一个损失函数的计算方式或者想把推理延迟从200ms压到50ms就完全不知道从哪里下手。后来我强迫自己用NumPy手写了一遍前向传播和反向传播用纯Python写了一遍数据加载器用Flask从零搭了一个推理接口才真正把“AI工程”这四个字吃透。这篇文章我就把这套从零构建的完整思路和实操细节拆开来讲包括每一步为什么这么做、有哪些坑、怎么验证自己做对了。2. 环境底座别急着装框架先把计算资源管明白2.1 硬件资源的真实盘点与显存预算从零做AI工程第一步不是pip install torch而是搞清楚你手里到底有多少计算资源。我见过太多人上来就加载一个7B参数的模型结果显存直接爆掉然后开始怀疑人生。这里有一个粗略的显存估算公式你在动手之前一定要先算一遍推理阶段显存占用 ≈ 参数量 × 精度字节数 × 1.2额外开销比如一个7B70亿参数的模型用FP16精度加载参数量占显存大约是 7 × 10^9 × 2 bytes ≈ 14GB再加上激活值、KV Cache、CUDA上下文等开销实际需要18GB到20GB。如果你只有一张16GB显存的卡那连加载都加载不进去。这时候你有几个选择用量化把精度降到INT8或INT4或者用CPU卸载部分层或者换更小的模型。但每一种选择都有代价量化会损失精度CPU卸载会增加延迟换小模型会降低效果。这些取舍就是AI工程的核心。训练阶段的显存估算更复杂除了模型参数还要存梯度、优化器状态、激活值。用Adam优化器做全量微调显存占用大概是推理的4到6倍。所以7B模型全量微调没有80GB显存的卡基本不用想。这也是为什么现在参数高效微调方法这么流行——不是大家不想全量微调是硬件不允许。提示在Linux环境下用nvidia-smi查看显存时注意Memory-Usage那一栏显示的是已分配显存但CUDA上下文本身就会占用几百MB实际可用显存比标称值少。做预算时留出至少10%的余量。2.2 Python环境隔离与依赖版本锁定AI工程最让人头疼的问题之一就是依赖冲突。你今天跑通的代码明天换台机器就报错大概率是某个包的版本变了。我的做法是每个项目一个独立的虚拟环境并且用pip freeze把精确版本号写进requirements.txt。不要用或~这种模糊版本AI领域的包更新极快一个小版本变动就可能改变默认行为。python -m venv ai-env source ai-env/bin/activate pip install numpy1.26.4 pip install torch2.2.0 --index-url https://download.pytorch.org/whl/cu121 pip freeze requirements.txt这里有个细节PyTorch的安装要指定CUDA版本对应的索引地址否则pip可能给你装一个CPU-only的版本跑起来慢得你想砸电脑。装完之后一定要验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果cuda.is_available()返回False先检查驱动版本和CUDA版本是否匹配再检查是不是装成了CPU版本。这个问题我至少遇到过五次每次都是因为pip源选错了。2.3 数据管道的I/O瓶颈预判很多人把注意力全放在模型上忽略了数据加载才是训练流程中最容易卡脖子的环节。我做过一个实验同样的模型和数据集用DataLoader默认参数和优化后的参数训练速度差了将近一倍。关键参数是num_workers和pin_memory。num_workers控制用几个子进程加载数据。设成0意味着在主进程里加载会阻塞训练设成CPU核心数通常是个好起点但也不是越大越好因为进程间通信本身有开销。pin_memoryTrue会把数据锁在页锁定内存里加速CPU到GPU的传输。还有一个容易被忽略的点如果你的数据集是小文件比如几KB的图片一定要提前打包成一个大文件如LMDB或WebDataset格式否则磁盘I/O会成为绝对瓶颈。from torch.utils.data import DataLoader loader DataLoader( dataset, batch_size32, shuffleTrue, num_workers8, pin_memoryTrue, prefetch_factor4, persistent_workersTrue )persistent_workersTrue这个参数在PyTorch 1.7之后才有作用是让worker进程在epoch之间不被销毁避免反复创建进程的开销。实测下来在小数据集上能省10%到15%的时间。3. 手写训练循环把反向传播从黑盒里拽出来3.1 用NumPy实现一个两层神经网络的前向与反向要真正理解AI工程我强烈建议你至少用NumPy手写一次完整的训练循环。不是让你以后都用NumPy做训练而是通过这个过程你会彻底明白loss.backward()到底干了什么。下面是一个两层全连接网络的核心代码import numpy as np # 初始化 W1 np.random.randn(784, 128) * 0.01 b1 np.zeros(128) W2 np.random.randn(128, 10) * 0.01 b2 np.zeros(10) def forward(x): z1 x W1 b1 a1 np.maximum(0, z1) # ReLU z2 a1 W2 b2 return z1, a1, z2 def backward(x, z1, a1, z2, y_true): batch x.shape[0] # softmax cross entropy 的梯度 probs np.exp(z2) / np.exp(z2).sum(axis1, keepdimsTrue) dz2 (probs - y_true) / batch dW2 a1.T dz2 db2 dz2.sum(axis0) da1 dz2 W2.T dz1 da1 * (z1 0) # ReLU导数 dW1 x.T dz1 db1 dz1.sum(axis0) return dW1, db1, dW2, db2这段代码里最关键的是dz2 (probs - y_true) / batch这一行。很多人不理解为什么softmax和交叉熵组合起来的梯度这么简洁其实推导一遍就明白了softmax的雅可比矩阵和交叉熵的梯度相乘恰好化简成了预测值减真实值。这就是数学的美妙之处也是框架帮你隐藏掉的东西。手写一遍之后你再去看PyTorch的autograd就会知道它本质上是在构建计算图然后沿着图反向传播链式法则。区别只是它支持任意复杂的网络结构并且用C做了高度优化。3.2 梯度检查确认你的反向传播没写错手写反向传播最大的风险是写错而且错了之后loss可能还在下降只是收敛得慢你很难发现。所以一定要做梯度检查。方法是用数值微分计算梯度的近似值和你解析求出的梯度对比。def numerical_gradient(f, x, eps1e-5): grad np.zeros_like(x) it np.nditer(x, flags[multi_index]) while not it.finished: idx it.multi_index old x[idx] x[idx] old eps fx1 f(x) x[idx] old - eps fx2 f(x) grad[idx] (fx1 - fx2) / (2 * eps) x[idx] old it.iternext() return grad相对误差控制在1e-7以内算正常1e-4到1e-7之间要警惕大于1e-4基本可以确定反向传播写错了。这个检查过程很慢但非常值得。我第一次手写的时候就是靠梯度检查发现ReLU的导数在z1等于0时处理错了。3.3 学习率与批大小的联动调参逻辑学习率和批大小不是独立超参数它们之间存在耦合关系。批大小增大k倍学习率通常也要相应增大但不是线性关系。经验法则是批大小翻倍学习率乘以sqrt(2)左右。更系统的做法是使用线性缩放规则或平方根缩放规则。批大小建议学习率范围适用场景321e-4 ~ 1e-3小数据集微调1283e-4 ~ 3e-3常规训练5121e-3 ~ 1e-2大数据集预训练20483e-3 ~ 3e-2大规模分布式还有一个实用技巧学习率预热。训练前几百步用很小的学习率然后线性增加到目标值再按余弦或阶梯衰减。这样做能显著提升训练稳定性尤其是Transformer类模型。我实测过一个文本分类任务加了预热之后前1000步的loss波动明显减小最终准确率也高了0.5个百分点。4. 推理服务化从模型文件到稳定接口的完整链路4.1 模型序列化与加载的性能陷阱训练完的模型要保存下来供推理使用。PyTorch提供了两种方式torch.save(model.state_dict())保存参数或者torch.save(model)保存整个模型。我强烈建议只保存state_dict因为保存整个模型会把类的定义也序列化进去换一个代码结构就加载不了而且文件体积更大。加载模型时有一个常见陷阱先加载到CPU再移到GPU比直接加载到GPU更省显存峰值。因为直接加载到GPU时PyTorch会先在GPU上分配内存再填充数据峰值显存是模型大小的两倍。正确做法是model MyModel() state_dict torch.load(model.pth, map_locationcpu) model.load_state_dict(state_dict) model model.half().cuda() # 转FP16再上GPU model.eval()model.eval()这行不能省它会关闭Dropout和BatchNorm的训练行为。我见过有人忘了加这行推理结果每次都不一样排查了半天才发现是Dropout在作怪。4.2 用Flask搭一个带批处理能力的推理接口从零搭推理服务Flask是最容易上手的起点。但直接写一个app.route然后同步调用模型吞吐量会非常低。核心优化点是动态批处理把短时间内到达的多个请求攒成一个批次一起推理充分利用GPU的并行能力。import threading import time from flask import Flask, request, jsonify app Flask(__name__) batch_buffer [] buffer_lock threading.Lock() BATCH_SIZE 8 MAX_WAIT 0.05 # 50ms def batch_worker(): while True: time.sleep(0.005) with buffer_lock: if len(batch_buffer) BATCH_SIZE or (batch_buffer and time.time() - batch_buffer[0][time] MAX_WAIT): batch batch_buffer[:BATCH_SIZE] del batch_buffer[:BATCH_SIZE] else: continue inputs [item[input] for item in batch] results model_inference(inputs) # 批量推理 for item, result in zip(batch, results): item[future][result] result item[future][event].set() threading.Thread(targetbatch_worker, daemonTrue).start()这个模式的核心思想是请求到达时不立即推理而是放入缓冲区由后台线程攒批后统一处理。每个请求用一个Event等待结果。实测下来在请求并发量中等的情况下吞吐量能提升3到5倍。代价是单个请求的延迟增加了最多50ms这个取舍要根据你的业务场景来定。4.3 超时控制与优雅降级推理服务最怕的是请求堆积导致雪崩。必须设置超时和队列上限。我的做法是每个请求设置一个最大等待时间超过就返回降级结果或错误码。同时请求队列设一个上限满了之后直接拒绝新请求而不是无限堆积。from concurrent.futures import ThreadPoolExecutor, TimeoutError executor ThreadPoolExecutor(max_workers4) app.route(/predict, methods[POST]) def predict(): data request.get_json() future executor.submit(process_request, data) try: result future.result(timeout2.0) # 2秒超时 return jsonify({result: result}) except TimeoutError: return jsonify({error: service busy, please retry}), 503这里用ThreadPoolExecutor而不是直接起线程是因为它自带队列管理和超时控制。max_workers设成GPU能同时处理的批次数设太大反而会导致显存竞争。注意Flask默认是单线程的生产环境要用gunicorn或uwsgi起多个worker。但多个worker意味着多个进程每个进程都会加载一份模型显存占用会成倍增加。所以更合理的架构是一个进程加载模型多个worker通过共享内存或RPC调用推理进程。这个改造稍微复杂一些但显存利用率会高很多。5. 性能压测与瓶颈定位用数据说话5.1 延迟分解从请求到响应的每一毫秒优化性能的第一步是知道时间花在哪里了。一个推理请求的延迟可以分解为网络传输、请求解析、数据预处理、模型推理、后处理、响应序列化。我习惯在代码里埋点记录每个阶段的时间戳。import time def predict_with_timing(data): t0 time.perf_counter() input_tensor preprocess(data) t1 time.perf_counter() with torch.no_grad(): output model(input_tensor) t2 time.perf_counter() result postprocess(output) t3 time.perf_counter() return { result: result, timing: { preprocess_ms: (t1 - t0) * 1000, inference_ms: (t2 - t1) * 1000, postprocess_ms: (t3 - t2) * 1000 } }实测数据通常显示小模型推理本身只占30%到50%的时间剩下全花在预处理和后处理上。比如图像任务里解码JPEG、缩放、归一化可能比模型前向还慢。这时候优化方向就不是换更快的模型而是用GPU做预处理或者用DALI这样的加速库。5.2 用wrk和py-spy做压力测试与火焰图压测工具我用得最多的是wrk轻量且能打出高并发。命令很简单wrk -t4 -c100 -d30s --latency -s post.lua http://localhost:5000/predict-t4是4个线程-c100是100个并发连接-d30s跑30秒--latency输出延迟分布。重点看P99延迟和每秒请求数。如果P99远高于P50说明有长尾请求通常是批处理等待或者GC导致的。定位CPU瓶颈用py-spy它能直接对运行中的Python进程采样并生成火焰图py-spy record -o profile.svg --pid 12345 --duration 30火焰图里横向越宽的调用栈占用CPU时间越多。我曾在火焰图里发现json.dumps占用了大量时间后来换成了orjson序列化耗时直接降了70%。5.3 常见瓶颈与对应解决策略瓶颈现象可能原因解决策略GPU利用率低数据加载慢增加num_workers预取数据显存溢出批大小过大减小批大小用梯度累积P99延迟高批处理等待减小MAX_WAIT或改用连续批处理吞吐上不去单进程限制多进程共享内存或换TritonCPU占用高预处理重GPU预处理或用C扩展这张表是我踩坑之后总结的基本上覆盖了80%的性能问题。每次遇到性能瓶颈先对照这张表排查能省很多时间。6. 从脚本到工程代码组织与可复现性6.1 配置管理别把超参数硬编码在代码里从零做AI工程很容易写成一个大脚本所有参数都写在代码里。这样做的后果是换个实验就要改代码改着改着就忘了哪个版本对应哪个结果。我的做法是用YAML配置文件管理所有超参数代码只读配置。# config/train.yaml model: hidden_size: 256 num_layers: 2 dropout: 0.1 training: batch_size: 64 learning_rate: 0.001 epochs: 50 warmup_steps: 500 data: train_path: data/train.csv val_split: 0.1然后用argparse或hydra加载配置。这样每次实验只需要换一个配置文件结果和配置一一对应复现起来非常方便。6.2 日志与实验追踪让每次运行都有迹可循日志不是print就完事了。我要求自己每个实验至少记录配置文件、git commit hash、开始和结束时间、每个epoch的loss和指标、最终模型路径。这些信息写到一个JSON文件里配合TensorBoard或WandB可视化。import json import hashlib from datetime import datetime def log_experiment(config, metrics, model_path): record { timestamp: datetime.now().isoformat(), config: config, config_hash: hashlib.md5(json.dumps(config, sort_keysTrue).encode()).hexdigest(), metrics: metrics, model_path: model_path } with open(experiments.jsonl, a) as f: f.write(json.dumps(record) \n)用JSONL格式追加写入不会覆盖之前的记录而且方便用pandas读取分析。config_hash能帮你快速找到相同配置的历史实验。6.3 单元测试保证数据处理和模型逻辑的正确性AI代码也需要单元测试尤其是数据处理部分。我至少会写三类测试输入输出形状测试、边界条件测试、数值稳定性测试。def test_preprocess_shape(): raw load_sample(test.jpg) tensor preprocess(raw) assert tensor.shape (3, 224, 224) assert tensor.dtype torch.float32 def test_preprocess_normalization(): raw np.ones((224, 224, 3), dtypenp.uint8) * 255 tensor preprocess(raw) assert tensor.max() 1.0 assert tensor.min() 0.0 def test_model_output_range(): model MyModel() x torch.randn(2, 3, 224, 224) with torch.no_grad(): out model(x) assert not torch.isnan(out).any() assert not torch.isinf(out).any()这些测试写起来很快但能帮你避免很多低级错误。比如有一次我改了归一化参数忘了同步更新测试结果推理结果全错靠单元测试才发现。7. 一些让我少走弯路的实操心得7.1 关于调试先跑通再优化别反过来我见过太多人一上来就追求极致性能用各种高级技巧结果代码跑都跑不通。正确的顺序是先用最笨的方法跑通全流程确认结果正确然后再逐步优化。比如推理服务先用Flask同步接口跑通再改批处理再加超时最后做多进程。每一步都验证正确性这样出问题的时候你知道是哪个环节引入的。7.2 关于显存及时释放和清空缓存PyTorch的显存管理有个特点删除了张量显存不一定立即归还给系统而是留在缓存池里供后续分配。这在循环推理时容易造成显存缓慢增长。我的习惯是在每个推理批次结束后手动清空缓存import torch import gc def cleanup(): gc.collect() torch.cuda.empty_cache()gc.collect()回收Python对象torch.cuda.empty_cache()释放未使用的显存缓存。注意不要在每个请求里都调用开销太大一般是在批次之间或者定时调用。7.3 关于版本锁定CUDA、驱动和框架的三角关系CUDA版本、显卡驱动版本、PyTorch版本三者必须匹配。我维护了一个对照表PyTorch版本推荐CUDA版本最低驱动版本2.0.x11.7 / 11.8450.80.022.1.x11.8 / 12.1450.80.022.2.x11.8 / 12.1450.80.022.3.x11.8 / 12.1450.80.02装之前先nvidia-smi看驱动版本再决定装哪个CUDA版本的PyTorch。如果驱动太老要么升级驱动要么降级PyTorch。这个三角关系搞不清楚环境问题能折腾你一整天。7.4 关于数据训练集和推理时的预处理必须完全一致这是最隐蔽的坑之一。训练时用了某种归一化方式推理时忘了同步模型效果直接崩掉。我的做法是把预处理逻辑封装成一个类或函数训练和推理共用同一份代码。并且把这个预处理类的参数保存到模型文件旁边加载模型时一起加载。class Preprocessor: def __init__(self, mean, std, size): self.mean mean self.std std self.size size def __call__(self, image): image resize(image, self.size) image (image - self.mean) / self.std return image # 保存 torch.save({ model_state: model.state_dict(), preprocessor: preprocessor }, checkpoint.pth)这样即使换了一台机器只要加载checkpoint预处理参数就是对的。7.5 关于监控上线只是开始不是结束推理服务上线之后必须持续监控几个核心指标QPS、P50/P99延迟、GPU利用率、显存占用、错误率。我用Prometheus加Grafana搭了一套简易监控每10秒采集一次数据。有一次发现P99延迟在每天下午三点准时飙升排查后发现是定时任务在那个时间点跑数据备份占用了磁盘I/O。如果没有监控这种问题根本发现不了。from prometheus_client import Counter, Histogram, start_http_server REQUEST_COUNT Counter(inference_requests_total, Total inference requests) REQUEST_LATENCY Histogram(inference_latency_seconds, Inference latency) REQUEST_LATENCY.time() def handle_request(data): REQUEST_COUNT.inc() return model_inference(data) start_http_server(8000)这套监控代码不到20行但带来的可见性是巨大的。强烈建议每个上线的推理服务都加上。8. 后续可以继续深挖的几个方向把上面这套从零构建的流程跑通之后你已经具备了AI工程的核心能力。接下来如果想继续深入有几个方向值得投入时间。第一是连续批处理比动态批处理更细粒度能在请求级别做调度vLLM和TensorRT-LLM都用了这个技术吞吐量能再上一个台阶。第二是模型量化从FP16到INT8再到INT4每一步都能显著降低显存和延迟但精度损失需要仔细评估。第三是分布式推理把模型切分到多张卡上用张量并行或流水线并行来跑更大的模型。第四是编译优化用Torch.compile或TensorRT把模型图编译成更高效的执行引擎实测能提升20%到50%的推理速度。这些方向每一个都够写一整篇文章但前提是你已经把基础链路走通了。如果连手写训练循环、搭推理服务、做性能压测这些基本功都没练扎实直接上这些高级技术只会让你在遇到问题时更加迷茫。从零开始一步一步来这才是ai-engineering-from-scratch真正的意义。
返回列表