导出模型为 ONNX
机器学习部署的“血泪”经验分享：从模型训练到服务上线那些事儿
这篇文章其实是源于我过去两年参与的一个图像识别项目。项目初期我们团队信心满满，模型在本地跑得飞起，在测试数据上准确率也漂亮得不像话，以为部署上线就是顺理成章的事儿。结果现实狠狠地给了我们一记响亮的耳光。
部署不是简单把模型 load 上去，然后扔到服务器就完事。它涉及到环境管理、性能调优、服务稳定性、监控报警等等多个环节。特别是在资源有限、高并发场景下，稍有不慎，系统就会崩溃或者响应迟缓，直接影响业务。
今天我就结合这个项目的实战经历，来聊聊我在机器学习部署过程中踩过的坑，以及一些我认为非常值得借鉴的最佳实践。
一、项目背景和挑战
我们的目标是为一个电商客户搭建一个商品图像识别系统，用于自动识别用户上传的商品图片并打标签（比如“T恤”、“牛仔裤”等），以提升推荐系统的准确性。
整个项目大致分为两个阶段：
- 模型开发阶段：我们基于 PyTorch 搭建了 ResNet-50 的迁移学习模型，使用了自建的商品图像数据集进行训练。
- 模型部署阶段：模型需要在 AWS EC2 实例上提供 API 接口，支持每天数十万次的请求访问。
听起来是不是挺标准的？但等到部署阶段，问题接踵而至。
二、遇到的主要问题与挑战
1. 环境不一致导致模型加载失败
我们在本地用的是 Python 3.8 + PyTorch 1.10，但是在云服务器上默认安装的是 Python 3.6 和 PyTorch 1.4。模型文件一加载就报错，提示某些新特性不支持。
🧨 小插曲：当时我们在凌晨上线的时候才发现这个问题，直接懵了两分钟才意识到是版本问题。
2. 模型预测速度太慢，QPS 上不去
一开始我们用 Flask 跑模型服务，单个请求处理时间平均 200ms，压测发现 QPS 勉强能到 15，完全扛不住预估的流量。
3. 显存占用太高，GPU 使用率不稳定
模型每次推理都要重新 load 到 GPU 上，而且 batch size 太小，显存利用率只有 30%，GPU 几乎处于空转状态。
4. 缺乏监控和服务治理机制
服务挂掉后没人知道，重启靠人工；日志杂乱无章；没有限流熔断，一炸全盘崩。
三、解决方案：一步步稳扎稳打
面对这些问题，我们没有慌，而是逐步优化架构、调整技术栈、完善运维流程。以下是我们的整体解决思路。
技术选型升级
| 模块 | 初始方案 | 最终方案 | 原因 |
|---|---|---|---|
| Web框架 | Flask | FastAPI + Gunicorn + Uvicorn Worker | 更好的异步支持，性能更好 |
| 模型服务化 | 单进程加载 | TorchServe（后来换成自研） | 提供更好的性能、批处理和版本控制 |
| 模型运行时 | CPU 推理 | GPU + ONNX Runtime | 利用 GPU 加速推理 |
| 服务部署方式 | 单节点 EC2 | Docker 容器 + Kubernetes 集群 | 支持弹性伸缩、滚动更新 |
| 日志监控 | 自定义 print | Prometheus + Grafana + ELK | 服务可视化，异常感知更快 |
架构图简化示意
[Client] → [Nginx/LB] → [Kubernetes Pods] → [ONNX + GPU 推理]
↘ [Prometheus/Grafana/ELK]
四、代码示例：如何做模型转换与高性能推理？
为了提高推理效率，我们最终将模型导出为 ONNX 格式，并使用 onnxruntime 进行推理。
import torch.onnx
model.eval()
dummy_input = torch.randn(1, 3, 224, 224) # 输入尺寸根据模型结构修改
torch.onnx.export(model, dummy_input, "resnet50.onnx", export_params=True)
# 使用 onnxruntime 推理
import numpy as np
import onnxruntime as ort
ort_session = ort.InferenceSession("resnet50.onnx")
outputs = ort_session.run(
None,
{'input': dummy_input.numpy()}
)
为了进一步提升吞吐量，我们将多个图片合并成 batch 进行批量推理，而不是单张处理。同时我们也利用了 CUDA 进行加速：
pip install onnxruntime-gpu # 一定要装带 GPU 的版本
五、部署过程中的几个关键细节
1. 批处理 vs 并发模型推理
我们尝试过两种方式：
- 同步逐条推理：每个请求进来自独立处理，缺点是吞吐低，GPU 无法有效利用。
- 队列式批处理：收集一定数量的请求组成 batch，一次性推理。这种方式可以显著提高 GPU 利用率，但也引入了延时。
我们最后选择了一个折中方案——设置最大等待时间 + 最小批大小，达到其中之一就触发推理。
from collections import deque
import threading
class BatchPredictor:
def __init__(self, model_path, max_batch_size=16, timeout=0.1):
self.model = load_model(model_path)
self.queue = deque()
self.lock = threading.Lock()
self.batching_thread = threading.Thread(target=self._process_loop)
self.batching_thread.daemon = True
self.batching_thread.start()
def _process_loop(self):
while True:
with self.lock:
if len(self.queue) < 1:
time.sleep(0.01)
continue
batch = list(self.queue)[:self.max_batch_size]
self.queue = deque(list(self.queue)[self.max_batch_size:])
# perform batch inference
predictions = self.model.predict(batch)
2. 模型热加载与服务降级机制
我们实现了一个简易的模型版本管理系统，允许通过 HTTP 触发模型 reload，这样可以在不停机的情况下完成模型更新。
同时在服务入口加了个简单的限流和熔断机制，防止雪崩效应。
from fastapi import FastAPI, HTTPException
from starlette.middleware import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
class RateLimitingMiddleware(BaseHTTPMiddleware):
def __init__(self, app, limit=100):
super().__init__(app)
self.limit = limit
self.requests = {}
async def dispatch(self, request, call_next):
client_ip = request.client.host
now = time.time()
self.requests.setdefault(client_ip, []).append(now)
self.requests[client_ip] = [t for t in self.requests[client_ip] if now - t < 60]
if len(self.requests[client_ip]) > self.limit:
raise HTTPException(status_code=429, detail="Too many requests")
return await call_next(request)
六、部署之后的效果
最终我们达成了如下效果：
| 指标 | 优化前 | 优化后 |
|---|---|---|
| QPS | 15 | 240+ |
| 平均延迟 | ~200ms | ~60ms |
| GPU 利用率 | 30% → 80% | |
| 故障恢复时间 | 依赖手动重启 | 自动重启 + 健康检查 |
| 支持多模型 | ❌ | ✅（支持灰度发布） |
我们还实现了自动扩缩容，Kubernetes 根据 CPU/GPU 使用率自动调整实例数，成本降低了不少。
七、踩坑总结 & 经验建议
1. 环境一致性才是王道
不要低估环境差异带来的影响。哪怕是一点 minor 版本的变化，都可能导致灾难。我推荐大家使用容器镜像打包整个环境，包括 Python、CUDA、CUDNN 和模型文件。例如用 Dockerfile：
FROM nvidia/cuda:11.7.1-base
RUN apt-get update && apt-get install -y python3-pip
COPY resnet50.onnx /models/
COPY requirements.txt .
RUN pip install -r requirements.txt
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "80"]
2. 别一开始就追求复杂的技术栈
TorchServe、TensorRT、Seldon、TF Serving 听起来都很高级，但在早期验证阶段完全没必要。先跑通再说，再根据实际情况迭代升级。
3. 写好文档，写好日志
别想着以后再补文档，否则你一定会后悔。部署服务时，务必记录每一步做了什么配置，用了哪些命令。日志也一样，别只输出 info 级别，debug、warn、error 要分清楚。
4. 测试！测试！测试！
线上出问题往往是因为没覆盖某个 edge case。一定要写好单元测试、集成测试，还要模拟各种异常情况（网络中断、磁盘满、GPU 异常等）。
5. 不要忽略业务端对接体验
接口设计要简洁清晰，返回格式统一，错误码明确，文档齐全。我们曾因为接口设计不合理导致客户端不断报错，浪费了很多时间排查。
八、最后想说的话
机器学习部署这件事儿，说难也难，说容易其实也有迹可循。关键是不要怕试错，也不要迷信大厂那一套架构，适合自己的才是最好的。
我个人最大的感悟是：一个 ML 工程师如果不了解部署和运维，就永远只是个“玩具开发者”。 只有当你把自己的模型真正放到生产环境，支撑起实际业务，才算是真正完成了闭环。
如果你正在准备或已经进入 ML 部署的阶段，我希望你看到这篇文章能少走些弯路，别像我当初那样被现实打得措手不及。
当然，也欢迎你在评论区留言交流你的经验和困惑。愿我们都能做出“跑得快、稳得住”的 AI 应用 💪
作者简介
一名曾在某电商平台负责 CV/NLP 相关算法工程化的全栈工程师，现专注于机器学习模型服务化与边缘部署方向。喜欢把理论落地，也热爱开源社区。欢迎关注我的 GitHub 或博客了解更多内容。
评论 0