AI 模型部署从入门到生产 —— ONNX 转换、TensorRT 加速、推理服务搭建
前言
模型训练出来只是第一步，让它稳定高效地服务线上请求才是真正的 challenge。这篇文章覆盖模型部署的完整流程——从模型导出到生产级推理服务。
Step 1：模型导出与格式转换
PyTorch → ONNX
ONNX（Open Neural Network Exchange）是模型部署的中间格式，几乎所有的推理框架和硬件都支持它。
importtorchimporttorch.onnx model=YourModel()model.eval()dummy_input=torch.randn(1,3,224,224)torch.onnx.export(model,dummy_input,"model.onnx",input_names=["input"],output_names=["output"],dynamic_axes={"input":{0:"batch_size"},"output":{0:"batch_size"},},opset_version=17,)关键参数：
- dynamic_axes：声明哪些维度是可变的，比如 batch_size 和序列长度
- opset_version：越高支持的算子越多，但硬件兼容性可能下降。推荐 17-19
ONNX 优化
导出的 ONNX 模型可以用onnxruntime做优化：
python-monnxruntime.tools.convert_onnx_models_to_ort model.onnx或者在代码中优化：
importonnxruntimeasortimportonnxfromonnxruntime.transformersimportoptimizer opt=optimizer.optimize_model("model.onnx",model_type="bert",# 根据模型类型选择num_heads=12,hidden_size=768,opt_level=99,# 最大优化)opt.save_model_to_file("model_optimized.onnx")优化效果：一般能提升 1.5x-3x 的推理速度。
检查 ONNX 模型的正确性
# 验证输出一致性withtorch.no_grad():torch_output=model(dummy_input).numpy()ort_session=ort.InferenceSession("model.onnx")ort_input={ort_session.get_inputs()[0].name:dummy_input.numpy()}ort_output=ort_session.run(None,ort_input)# 检查误差diff=np.abs(torch_output-ort_output[0]).max()print(f"Max diff:{diff}")# 对于 FP32，diff < 1e-5 是正常的# 对于 FP16/INT8，diff < 1e-2 是可接受的Step 2：推理加速 —— TensorRT
TensorRT 是 NVIDIA 的推理优化引擎，可以在 GPU 上达到极致性能。
ONNX → TensorRT
importtensorrtastrt logger=trt.Logger(trt.Logger.WARNING)builder=trt.Builder(logger)network=builder.create_network(1<<int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser=trt.OnnxParser(network,logger)withopen("model.onnx","rb")asf:parser.parse(f.read())config=builder.create_builder_config()config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,1<<30)# 1GB# FP16 精度（精度几乎无损，速度提升 2x）ifbuilder.platform_has_fast_fp16:config.set_flag(trt.BuilderFlag.FP16)# INT8 量化（需要校准数据集，速度再提升 2x）# config.set_flag(trt.BuilderFlag.INT8)engine=builder.build_serialized_network(network,config)withopen("model.trt","wb")asf:f.write(engine)精度与速度权衡
| 精度 | 推理速度 | 模型大小 | 质量影响 | 适用场景 |
|---|---|---|---|---|
| FP32 | 1x | 100% | 无 | 基准、调试 |
| FP16 | 1.5-2x | 50% | 几乎无 | 生产推荐 |
| INT8 | 3-4x | 25% | 轻微 | 对延迟极度敏感 |
| INT4 | 5-6x | 12.5% | 明显 | 边缘设备 |
通用建议：FP16 是性价比最高的选择。INT8 需要校准数据集，且对某些模型有精度风险。
TensorRT 推理
importtensorrtastrtimportpycuda.driverascudaclassTRTInference:def__init__(self,engine_path):logger=trt.Logger(trt.Logger.WARNING)withopen(engine_path,"rb")asf:self.engine=trt.Runtime(logger).deserialize_cuda_engine(f.read())self.context=self.engine.create_execution_context()# 分配 GPU 内存self.inputs=[]self.outputs=[]foriinrange(self.engine.num_io_tensors):name=self.engine.get_tensor_name(i)shape=self.engine.get_tensor_shape(name)size=trt.volume(shape)dtype=trt.nptype(self.engine.get_tensor_dtype(name))# 分配内存...# 完整代码见注definfer(self,input_numpy):self.inputs[0].host=input_numpy.astype(np.float32)[cuda.memcpy_htod(inp.device,inp.host)forinpinself.inputs]self.context.execute_v2([inp.deviceforinpinself.inputs+self.outputs])[cuda.memcpy_dtoh(out.host,out.device)foroutinself.outputs]returnself.outputs[0].hostStep 3：推理服务框架
三种主流方案对比
| 框架 | 适用场景 | 特点 |
|---|---|---|
| Triton Inference Server | 多模型混合、A/B 测试 | NVIDIA 官方，支持多种硬件和框架 |
| vLLM | LLM 推理（7B-70B） | PagedAttention，连续批处理 |
| SGLang | LLM 推理（结构化输出） | RadixAttention，JSON Mode |
用 vLLM 部署 LLM
fromvllmimportLLM,SamplingParams llm=LLM(model="meta-llama/Llama-3.2-3B",tensor_parallel_size=1,# 单卡部署gpu_memory_utilization=0.9,# GPU 显存利用率max_num_seqs=256,# 最大并发序列数enable_prefix_caching=True,# 前缀缓存加速)sampling_params=SamplingParams(temperature=0.7,top_p=0.9,max_tokens=1024,)outputs=llm.generate(["请解释什么是 Transformer"],sampling_params)print(outputs[0].outputs[0].text)生产部署用 OpenAI 兼容 API：
python-mvllm.entrypoints.openai.api_server--modelmeta-llama/Llama-3.2-3B--port8000--gpu-memory-utilization0.9然后直接当 OpenAI API 用：
fromopenaiimportOpenAI client=OpenAI(base_url="http://localhost:8000/v1",api_key="sk-xxx",# vLLM 不验证 key，随便填)response=client.chat.completions.create(model="meta-llama/Llama-3.2-3B",messages=[{"role":"user","content":"Hello"}],)vLLM 关键配置
LLM(model="...",# 显存配置gpu_memory_utilization=0.85-0.95,# 留一些给 CUDA kernelmax_model_len=8192,# 最大上下文长度# 性能配置num_scheduler_steps=8,# 调度步数，越大吞吐越高enable_chunked_prefill=True,# 分块 prefill，降低 TTFT# 量化quantization="awq",# AWQ 量化（需要先量化模型）# 分布式tensor_parallel_size=1,# >1 时需多卡)Step 4：性能压测与监控
关键指标
| 指标 | 英文 | 含义 | 好值 |
|---|---|---|---|
| 首 Token 延迟 | TTFT | 从请求到第一个 Token 的时间 | < 500ms |
| Token 生成速度 | TPOT | 每个 Token 的生成时间 | < 50ms |
| 吞吐量 | Throughput | 每秒生成的 Token 数 | > 1000 tokens/s |
| 并发数 | Concurrency | 同时处理的请求数 | 视场景而定 |
压测工具
# 用 vllm 自带的 benchmarkpython-mvllm.benchmarks.benchmark_throughput--modelmeta-llama/Llama-3.2-3B--datasetShareGPT_V3_unfiltered_cleaned_split.json --num-prompts1000# 用 locust 做 HTTP 压测pipinstalllocust locust--hosthttp://localhost:8000生产级部署架构
┌─────────────┐ │ Load │ │ Balancer │ └──────┬──────┘ │ ┌───────────────┼───────────────┐ │ │ │ ┌─────▼─────┐ ┌─────▼─────┐ ┌─────▼─────┐ │ Worker 1 │ │ Worker 2 │ │ Worker 3 │ │ vLLM/Triton│ │ vLLM/Triton│ │ vLLM/Triton│ └───────────┘ └───────────┘ └───────────┘ │ │ │ ┌─────▼───────────────▼───────────────▼─────┐ │ 共享显存 / 模型分片 │ └───────────────────────────────────────────┘- 水平扩展：多个 Worker 加 Load Balancer
- 动态缩扩容：根据队列长度自动调整 Worker 数量
- 缓存层：Redis 缓存常见问题的回答（适合 LLM 场景，回答可以复用）
总结
模型部署的核心链路：
PyTorch → ONNX → TensorRT → Triton/vLLM → Load Balancer → API ① ② ③ ④ ⑤- PyTorch 导出 ONNX（设置 dynamic_axes）
- ONNX 优化（onnxruntime transformers）
- TensorRT 转换（FP16 是最优性价比）
- 推理框架部署（LLM 用 vLLM，多模态用 Triton）
- 负载均衡与缓存
每个环节都有优化空间，但先跑通全链路，再逐步压测优化是最务实的方式。
觉得有用？点赞 + 收藏 + 关注。