如何部署pytorch模型

部署 PyTorch 模型通常涉及模型导出 → 选择部署方式 → 服务化/集成 → 优化与监控几个阶段。下面我按常见场景给你一个清晰、可落地的指南。

一、模型准备（训练完成后）

1. 保存模型

常见两种方式：

方式 A：保存整个模型（不推荐生产）

torch.save(model, "model.pth")

方式 B：只保存权重（推荐）

torch.save(model.state_dict(), "model_weights.pth")

部署时通常配合：

model = MyModel()
model.load_state_dict(torch.load("model_weights.pth"))
model.eval()

二、模型导出（非常重要）

✅ 推荐：导出为 TorchScript

example_input = torch.rand(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model.pt")

或（支持控制流）：

scripted_model = torch.jit.script(model)
scripted_model.save("model.pt")

✅ 优点：

• 不依赖 Python 训练环境
• 可在 C++ / 服务端运行

三、常见部署方式（按使用场景）

✅ 场景 1：Python Web 服务（最常见）

技术栈：Flask / FastAPI + PyTorch

from fastapi import FastAPI
import torch

app = FastAPI()
model = torch.jit.load("model.pt")
model.eval()

@app.post("/predict")
def predict(data: list):
    x = torch.tensor(data)
    with torch.no_grad():
        y = model(x)
    return y.tolist()

启动：

uvicorn main:app --host 0.0.0.0 --port 8000

✅ 场景 2：高性能推理（生产推荐）

TorchServe（官方）

pip install torchserve torch-model-archiver

打包：

torch-model-archiver \
  --model-name mymodel \
  --serialized-file model.pt \
  --export-path model_store \
  --runtime python

启动：

torchserve --start --model-store model_store --models mymodel

✅ 场景 3：移动端 / 边缘设备

PyTorch → ONNX → 移动端

torch.onnx.export(
    model,
    example_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"]
)

然后使用：

• ONNX Runtime Mobile
• TensorRT（NVIDIA）
• NCNN / MNN（Android/iOS）

✅ 场景 4：C++ 部署

#include <torch/script.h>

auto model = torch::jit::load("model.pt");
auto input = torch::rand({1, 3, 224, 224});
auto output = model.forward({input}).toTensor();

四、性能优化建议

• ✅ model.eval()
• ✅ with torch.no_grad():
• ✅ 使用 batch 推理
• ✅ GPU：model.cuda() + input.cuda()
• ✅ 使用 FP16 / INT8
• ✅ TensorRT / ONNX Runtime 加速

五、简单部署架构示例

客户端
  ↓
API (FastAPI / TorchServe)
  ↓
PyTorch 推理
  ↓
返回结果

六、你可以告诉我：

• ✅ 模型类型（CV / NLP / 推荐）
• ✅ 部署环境（服务器 / 云 / 手机 / 嵌入式）
• ✅ 是否有 GPU
• ✅ 实时性要求

我可以直接给你可运行的部署代码模板。