PyTorch在Ubuntu上如何进行模型部署与监控

Ubuntu上PyTorch模型的部署与监控实践指南

一 环境准备与模型导出

• 建议使用Ubuntu 20.04/22.04 LTS，先创建隔离环境（conda/venv），再安装与CUDA匹配的PyTorch：
  • CPU版：pip install torch torchvision torchaudio
  • GPU版（示例CUDA 11.7）：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
• 验证环境：python -c "import torch; print(torch.__version__, torch.cuda.is_available())"
• 模型导出为可部署格式（二选一或同时保留）：
  • TorchScript：scripted = torch.jit.script(model); scripted.save("model.pt")
  • ONNX：准备示例输入dummy_input，执行torch.onnx.export(model, dummy_input, "model.onnx")
• 可选优化（提升吞吐/降低延迟，上线前需评估精度回归）：
  • 混合精度：torch.cuda.amp.autocast() + GradScaler
  • 量化：torch.quantization（动态/静态）
  • 剪枝：torch.nn.utils.prune

二 部署路径与示例

• 方案A 原生Python服务（Flask/FastAPI）
  • 适用：中小流量、快速上线
  • 要点：模型加载到内存、设置model.eval()、预处理/后处理、用Gunicorn/Uvicorn多进程/异步运行
  • 示例（FastAPI）：
    • 安装：pip install fastapi uvicorn[standard]
    • 服务代码要点：
      • model = MyModel(); model.load_state_dict(torch.load("model.pth", map_location="cpu")); model.eval()
      • 推理用with torch.no_grad():
    • 启动：uvicorn app:app --host 0.0.0.0 --port 5000
• 方案B TorchServe（官方模型服务框架）
  • 适用：生产级、多模型管理、自动批处理、可观测性强
  • 快速步骤：
    • 安装：pip install torchserve torch-model-archiver torch-workflow-archiver
    • 打包：torch-model-archiver --model-name mnist --version 1.0 --model-file mnist.py --serialized-file mnist_cnn.pt --handler mnist_handler.py
    • 启动容器（宿主机端口3000推理、3001管理）：
      • docker run --rm -it -p 3000:8080 -p 3001:8081 pytorch/torchserve
• 方案C ONNX Runtime（跨平台高性能推理）
  • 适用：追求极致CPU/GPU推理性能、跨框架部署
  • 要点：用onnxruntime加载model.onnx推理，结合异步框架或批量队列提升吞吐
• 可选容器化
  • Dockerfile示例：
    • FROM python:3.10-slim
    • WORKDIR /app
    • COPY requirements.txt .
    • RUN pip install -r requirements.txt
    • COPY . .
    • CMD ["gunicorn", "app:app", "-b", "0.0.0.0:5000", "--timeout", "120"]
  • 构建与运行：docker build -t my-pytorch-app . && docker run -p 5000:5000 my-pytorch-app

三 运行期监控与可观测性

• 系统资源
  • GPU：watch -n 1 nvidia-smi（显存、利用率、温度、功耗）
  • CPU/内存：htop 或 top
  • 进程内采集：psutil 记录CPU%、RSS内存等，便于写入日志/打点
• 训练/推理指标与可视化
  • TensorBoard：pip install tensorboard torch.utils.tensorboard
    • 记录：SummaryWriter(log_dir="runs").add_scalar("Loss/train", loss, step)
    • 查看：tensorboard --logdir=runs
  • 模型结构与图可视化
    • Netron：pip install netron；netron model.onnx 在浏览器查看网络拓扑与维度
    • Torchinfo：summary(model, input_size=(B, C, H, W)) 查看参数量与层信息
• 性能剖析
  • PyTorch Profiler：定位CPU/GPU热点与内存瓶颈
    • 示例：
      • with torch.profiler.profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True, profile_memory=True) as prof:
      • outputs = model(inputs)
      • print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

四 上线检查清单与常见问题

• 上线前
  • 基准测试：固定batch与输入形状，测吞吐（QPS）、P95/P99延迟、显存/内存占用
  • 精度回归：对比FP32与量化/剪枝后模型的验证集指标
  • 资源与限流：设置合理的worker数与队列，防止OOM与雪崩
  • 日志与告警：记录请求耗时、异常堆栈、GPU温度/显存阈值告警
• 常见问题
  • CUDA OOM：减小batch、启用梯度检查点/CPU卸载、开启TorchServe动态批处理
  • 精度下降：量化/剪枝后做全量评估，必要时回退或做蒸馏
  • 依赖冲突：使用虚拟环境/容器固化依赖版本
  • 端口占用/跨域：确认宿主机端口映射与反向代理CORS配置正确