Ubuntu下PyTorch模型如何部署

Ubuntu下PyTorch模型部署全流程

一 环境准备与GPU验证

• 更新系统并安装基础工具：sudo apt update && sudo apt upgrade -y && sudo apt install -y python3 python3-pip
• 建议使用虚拟环境隔离依赖：python3 -m venv venv && source venv/bin/activate
• 安装与硬件匹配的PyTorch（CPU或CUDA）。例如：pip install torch torchvision torchaudio（CPU）；如需GPU，请选择与驱动匹配的CUDA版本，如：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
• 验证安装与设备：python - <<‘PY’ import torch print(“torch:”, torch.version, “cuda:”, torch.cuda.is_available(), “device:”, torch.cuda.get_device_name(0) if torch.cuda.is_available() else “CPU”) PY
• 如启用GPU，确保已安装对应版本的NVIDIA驱动、CUDA与cuDNN，并使用nvidia-smi监控资源与版本匹配情况。

二 模型导出与格式选择

• 保存与加载的常用方式
  • 仅保存权重（推荐）：torch.save(model.state_dict(), “model.pth”)；加载：model.load_state_dict(torch.load(“model.pth”))
  • 保存整个模型：torch.save(model, “model.pt”)；加载：model = torch.load(“model.pt”)
  • 推理前务必调用：model.eval()
• 导出为TorchScript（便于跨环境、C++部署）
  • 示例：import torch; model.eval(); example = torch.randn(1,3,224,224); traced = torch.jit.trace(model, example); traced.save(“model.pt”)
• 导出为ONNX（便于跨框架与多引擎推理）
  • 示例：torch.onnx.export(model, example, “model.onnx”, opset_version=10, input_names=[“input”], output_names=[“output”], dynamic_axes={“input”:{0:“batch”}, “output”:{0:“batch”}})
  • 推理引擎可选ONNX Runtime：pip install onnxruntime；示例：ort_session = onnxruntime.InferenceSession(“model.onnx”); ort_inputs = {“input”: input_np}; ort_outs = ort_session.run(None, ort_inputs)
• 导出要点
  • 动态批处理可用dynamic_axes；自定义算子需确认目标引擎支持或在导出时处理
  • 推理前统一预处理/后处理与训练保持一致（归一化、尺寸、类型）。

三 部署路径与示例

• 原生Python服务（Flask/FastAPI）
  • 特点：开发快、灵活；适合中小流量或内部服务
  • 示例（FastAPI）：
    • pip install fastapi uvicorn[standard]
    • 代码示例：
      • from fastapi import FastAPI
      • import torch, json
      • app = FastAPI()
      • model = torch.jit.load(“model.pt”, map_location=“cpu”).eval()
      • @app.post(“/predict”)
      • def predict(data: dict):
        • x = torch.tensor(data[“input”], dtype=torch.float32).unsqueeze(0)
        • with torch.no_grad(): y = model(x).squeeze().tolist()
        • return {“output”: y}
      • 运行：uvicorn app:app --host 0.0.0.0 --port 5000
• TorchServe（官方模型服务框架）
  • 特点：模型管理、版本化、批处理、并发与可观测性强，适合生产
  • Docker方式（CPU示例）：
    • docker pull pytorch/torchserve:latest
    • 准备模型目录：mkdir -p model_store && torch-model-archiver --model-name mymodel --version 1.0 --serialized-file model.pt --handler image_classifier --export-path model_store
    • 启动：docker run --rm -p 8080:8080 -p 8081:8081 -p 8082:8082 -v $(pwd)/model_store:/home/model-server/model-store pytorch/torchserve:latest
    • 访问：http://localhost:8080/predictions/mymodel
• ONNX Runtime高性能推理
  • 特点：跨平台、多语言、推理引擎优化成熟；适合高吞吐场景
  • 示例见“二”中ONNX Runtime代码；可结合线程池/异步IO与批量合并优化吞吐。

四 性能优化与常见问题

• GPU推理优化
  • 批处理与数据预取；固定输入尺寸减少开销
  • 混合精度：from torch.cuda.amp import autocast; with autocast(): output = model(input)
  • 启用cuDNN加速：torch.backends.cudnn.benchmark = True（输入尺寸固定时收益明显）
  • 监控与排障：nvidia-smi观察显存/利用率；CUDA out of memory时减小batch或改用amp
• TensorRT加速（NVIDIA GPU）
  • 将模型导出为ONNX后，使用TensorRT进行图优化与内核选择，可显著提升吞吐与降低延迟（适合高并发在线服务）
• 常见问题速查
  • torch.cuda.is_available() 为False：检查驱动/CUDA/库版本匹配与容器设备映射
  • 导出/推理结果不一致：确认model.eval()、预处理/后处理一致、动态轴设置正确
  • 服务并发瓶颈：增加Worker数量、启用批量、使用异步框架或切换到更高性能引擎（如TensorRT/ONNX Runtime）。

五 边缘与跨平台部署

• 移动端/嵌入式（NCNN）
  • 路径：PyTorch → TorchScript → PNNX → NCNN。示例：先用torch.jit.trace导出.pt，再用PNNX转换生成.ncnn.bin/.param，部署到ncnn推理引擎（支持C++/Android）
• ARM嵌入式（Rockchip RK3588）
  • 路径：PyTorch(.pth) → ONNX → RKNN Toolkit2 → RKNN模型 → 板端rknn-toolkit2-lite推理（Python/C）。