在Ubuntu上部署PyTorch模型的最佳实践

在Ubuntu上部署PyTorch模型可遵循以下最佳实践，涵盖环境搭建、模型转换、部署及优化等关键环节：

一、环境搭建

1. 安装依赖
   • 使用apt安装系统依赖：sudo apt update && sudo apt install python3 python3-pip build-essential。
   • 推荐通过conda管理虚拟环境：

     wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
     bash Miniconda3-latest-Linux-x86_64.sh
     conda create -n pytorch_env python=3.8 -y
     conda activate pytorch_env
     

2. 安装PyTorch
   • 根据CUDA版本选择安装命令（以CUDA 11.8为例）：

     conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia
     

     或使用pip（需指定CUDA版本）：

     pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
     

   • 验证安装：

     import torch
     print(torch.__version__, torch.cuda.is_available())
     

二、模型准备与转换

1. 模型保存
   • 训练完成后，将模型保存为.pt或.pth文件：

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

2. 格式转换（可选）
   • TorchScript：适用于非Python环境，支持跨平台部署：

     scripted_model = torch.jit.script(model)
     scripted_model.save('model_scripted.pt')
     

   • ONNX：支持跨框架部署（如TensorRT优化）：

     torch.onnx.export(model, input_tensor, 'model.onnx', opset_version=11)
     

三、部署方案

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)
  

• Flask/FastAPI Web服务
  • 安装框架：pip install flask。
  • 示例代码（Flask）：

    from flask import Flask, request, jsonify
    import torch
    
    app = Flask(__name__)
    model = torch.load('model.pth', map_location=torch.device('cpu'))
    model.eval()
    
    @app.route('/predict', methods=['POST'])
    def predict():
        data = request.json['input']
        input_tensor = torch.tensor(data).unsqueeze(0)
        with torch.no_grad():
            output = model(input_tensor)
        return jsonify(output.tolist())
    
    if __name__ == '__main__':
        app.run(host='0.0.0.0', port=5000)
    

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

2. 生产环境部署

• Gunicorn + Flask
  提升Web服务并发能力：

  pip install gunicorn
  gunicorn -w 4 -b 0.0.0.0:5000 app:app
  

• Docker容器化
  • 创建Dockerfile：

    FROM python:3.8-slim
    WORKDIR /app
    COPY requirements.txt .
    RUN pip install -r requirements.txt
    COPY . .
    CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app:app"]
    

  • 构建并运行：

    docker build -t pytorch-model-server .
    docker run -p 5000:5000 pytorch-model-server
    

• TorchServe（官方推荐）
  适合大规模部署，支持模型热更新和监控：

  pip install torchserve torch-model-archiver
  torch-model-archiver --model-name my_model --version 1.0 --serialized-file model.pth --handler image_classifier
  torchserve --start --model-store ./ --models my_model.mar
  

四、性能优化

1. GPU加速
   • 确保CUDA和cuDNN已正确安装，通过torch.cuda.is_available()验证。
   • 在代码中指定GPU设备：

     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model.to(device)
     input_tensor = input_tensor.to(device)
     

2. 批量推理
   • 合并多个输入样本进行批量处理，提升吞吐量：

     batch_input = torch.cat([input1, input2, input3], dim=0)
     with torch.no_grad():
         batch_output = model(batch_input)
     

五、安全与维护

• 依赖管理：使用requirements.txt记录依赖版本，避免环境冲突：

  pip freeze > requirements.txt
  

• 日志与监控：通过logging模块记录服务日志，结合Prometheus+Grafana监控模型性能。
• 定期更新：定期升级PyTorch版本及依赖库，修复安全漏洞。

参考资料