PyTorch在Ubuntu上的模型部署方法

以下是在Ubuntu上部署PyTorch模型的主要方法，步骤基于搜索资源整合：

一、环境准备

1. 安装系统依赖
   • 更新系统：sudo apt update && sudo apt upgrade
   • 安装NVIDIA驱动、CUDA和cuDNN（需匹配PyTorch版本）。
2. 创建虚拟环境
   • 使用Conda或venv隔离环境：

     # Conda  
     conda create -n pytorch_env python=3.8  
     conda activate pytorch_env  
     # 或venv  
     python3 -m venv pytorch_env  
     source pytorch_env/bin/activate  
     

3. 安装PyTorch
   • CPU版本：pip install torch torchvision torchaudio
   • GPU版本：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113（需替换为实际CUDA版本）。

二、模型准备

• 训练并保存模型为.pt或.pth文件，例如：

  torch.save(model.state_dict(), 'model.pth')  
  

三、部署方式

1. 直接运行脚本（轻量级）

• 编写Python脚本加载模型并处理输入：

  import torch  
  from model import MyModel  # 假设模型定义在此文件中  
  
  model = MyModel()  
  model.load_state_dict(torch.load('model.pth'))  
  model.eval()  
  
  # 示例输入（需根据模型调整）  
  input_data = torch.randn(1, 3, 224, 224)  
  with torch.no_grad():  
      output = model(input_data)  
  print(output)  
  

  运行：python deploy.py。

2. Web服务部署（Flask/FastAPI）

• Flask示例：
  1. 安装Flask：pip install flask
  2. 创建app.py：

     from flask import Flask, request, jsonify  
     import torch  
     from model import MyModel  
     
     app = Flask(__name__)  
     model = MyModel()  
     model.load_state_dict(torch.load('model.pth'))  
     model.eval()  
     
     @app.route('/predict', methods=['POST'])  
     def predict():  
         data = request.json['input']  # 假设输入为JSON格式  
         input_tensor = torch.tensor(data).unsqueeze(0)  
         with torch.no_grad():  
             output = model(input_tensor)  
         return jsonify({'output': output.tolist()})  
     
     if __name__ == '__main__':  
         app.run(host='0.0.0.0', port=5000)  
     

  3. 运行服务：python app.py，通过HTTP请求调用。

3. 生产环境部署

• Gunicorn + Flask：
  安装Gunicorn：pip install gunicorn，运行：
  gunicorn -w 4 -b 0.0.0.0:5000 app:app（-w指定进程数）。
• Docker容器化：
  1. 编写Dockerfile：

     FROM pytorch/pytorch:latest  
     COPY . /app  
     WORKDIR /app  
     RUN pip install -r requirements.txt  
     CMD ["python", "app.py"]  
     

  2. 构建并运行：

     docker build -t pytorch-model .  
     docker run -p 5000:5000 pytorch-model  
     ```。  
     
     

4. 高性能部署（TorchServe/ONNX）

• TorchServe（官方推荐）：
  1. 安装TorchServe：pip install torchserve torch-model-archiver
  2. 将模型转换为TorchScript格式：

     scripted_model = torch.jit.script(model)  
     scripted_model.save('model.ts')  
     

  3. 启动服务：
     torchserve --start --model_name=my_model --model_path=/path/to/model.ts。
• ONNX转换：
  将模型导出为ONNX格式，部署到支持ONNX的推理引擎（如TensorRT）：
  torch.onnx.export(model, input, 'model.onnx')。

四、验证与监控

• 通过curl或Postman测试API接口：

  curl -X POST -H "Content-Type: application/json" -d '{"input": [1,2,3]}' http://localhost:5000/predict  
  

• 生产环境中建议使用Nginx反向代理、Prometheus监控服务状态。

注意事项

• GPU部署需确保CUDA/cuDNN版本与PyTorch兼容。
• 大规模部署建议使用Docker + Kubernetes实现弹性扩缩容。

参考来源：