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）。

最新问答