Ubuntu上PyTorch如何进行模型部署

Ubuntu上PyTorch模型部署完整流程

1. 环境准备：安装PyTorch及依赖

在Ubuntu上部署PyTorch模型前，需先配置基础环境。建议使用虚拟环境（如conda或venv）隔离依赖，避免冲突。

• 安装Python基础工具：通过apt安装Python3及pip包管理器。

  sudo apt update &
      &
   sudo apt install -y python3 python3-pip python3-venv
  

• 创建并激活虚拟环境：

  python3 -m venv pytorch_deploy_env
  source pytorch_deploy_env/bin/activate
  

• 安装PyTorch：根据是否需要GPU加速，选择对应安装命令（以CUDA 11.8为例，需提前安装NVIDIA驱动和CUDA Toolkit）：

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

  若无需GPU，可直接安装CPU版本：

  pip install torch torchvision torchaudio
  

• 验证安装：运行Python脚本检查PyTorch版本及CUDA可用性。

  import torch
  print(torch.__version__)  # 输出PyTorch版本
  print(torch.cuda.is_available())  # 输出True表示CUDA可用
  

2. 模型准备：保存与加载

模型部署前需完成训练并保存，支持两种格式：完整模型（包含架构与权重）或状态字典（仅权重，推荐）。

• 保存模型：
  • 完整模型（直接保存整个模型对象）：

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

  • 状态字典（推荐，灵活性更高）：

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

• 加载模型：
  • 完整模型：直接加载即可使用（需确保PyTorch版本一致）。

    model = torch.load('full_model.pth')
    

  • 状态字典：需先实例化模型类，再加载权重。

    from model import MyModel  # 导入模型定义
    model = MyModel()
    model.load_state_dict(torch.load('model_state_dict.pth'))
    model.eval()  # 切换至评估模式（关闭Dropout/BatchNorm的随机性）
    

3. 推理脚本编写：基础预测

编写Python脚本实现模型加载与推理，支持批处理和预处理（如图像分类任务）。

import torch
from torchvision import transforms
from PIL import Image
from model import MyModel  # 替换为你的模型类

# 初始化模型
model = MyModel()
model.load_state_dict(torch.load('model_state_dict.pth'))
model.eval()

# 定义预处理（以ResNet为例）
preprocess = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像尺寸
    transforms.ToTensor(),          # 转换为Tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 推理函数
def predict(image_path):
    # 加载图像并预处理
    image = Image.open(image_path).convert('RGB')
    input_tensor = preprocess(image).unsqueeze(0)  # 添加batch维度
    
    # 推理（禁用梯度计算，提升速度）
    with torch.no_grad():
        output = model(input_tensor)
    
    # 后处理（如分类任务取最大概率类别）
    _, predicted_class = torch.max(output, 1)
    return predicted_class.item()

# 示例调用
if __name__ == "__main__":
    result = predict('test_image.jpg')
    print(f"Predicted class: {
result}
")

4. 生产环境部署：Web服务化

若需将模型作为API服务对外提供，可使用Flask（轻量级）或FastAPI（高性能）框架。以下以Flask为例：

• 安装Flask：

  pip install flask
  

• 创建Flask应用（app.py）：

  from flask import Flask, request, jsonify
  import torch
  from PIL import Image
  from io import BytesIO
  from model import MyModel
  import torchvision.transforms as transforms
  
  app = Flask(__name__)
  
  # 初始化模型
  model = MyModel()
  model.load_state_dict(torch.load('model_state_dict.pth'))
  model.eval()
  
  # 预处理函数
  preprocess = transforms.Compose([
      transforms.Resize((224, 224)),
      transforms.ToTensor(),
      transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  ])
  
  # 定义预测接口
  @app.route('/predict', methods=['POST'])
  def predict():
      if 'image' not in request.files:
          return jsonify({
  'error': 'No image uploaded'}
  ), 400
      
      # 读取图像文件
      file = request.files['image']
      image = Image.open(BytesIO(file.read())).convert('RGB')
      
      # 预处理与推理
      input_tensor = preprocess(image).unsqueeze(0)
      with torch.no_grad():
          output = model(input_tensor)
      
      # 返回结果
      _, predicted_class = torch.max(output, 1)
      return jsonify({
  'predicted_class': int(predicted_class.item())}
      )
  
  if __name__ == '__main__':
      app.run(host='0.0.0.0', port=5000, debug=False)  # 生产环境关闭debug模式
  

5. 生产环境部署：性能优化

为提升模型服务的性能与稳定性，需进行以下优化：

• 使用Gunicorn作为WSGI服务器：替代Flask自带的开发服务器，支持多进程并发。

  pip install gunicorn
  gunicorn -w 4 -b 0.0.0.0:5000 app:app  # 4个工作进程
  

• 使用Docker容器化：将应用与依赖打包成镜像，便于跨环境部署。
  • 创建Dockerfile：

    FROM python:3.8-slim
    WORKDIR /app
    COPY requirements.txt .  # 依赖文件
    RUN pip install --no-cache-dir -r requirements.txt
    COPY . .                # 复制应用代码
    CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app:app"]
    

  • 构建与运行镜像：

    docker build -t pytorch-model-server .
    docker run -d -p 5000:5000 --gpus all pytorch-model-server  # 若需GPU支持，添加--gpus all
    

• 启用GPU加速：若使用CUDA模型，需安装NVIDIA Container Toolkit并添加--gpus all参数运行Docker容器。

6. 高级部署：TorchServe（官方工具）

TorchServe是PyTorch官方提供的模型服务框架，支持REST API、gRPC和模型热加载，适合大规模生产部署。

• 安装TorchServe：

  pip install torchserve torch-model-archiver
  

• 打包模型：将模型转换为TorchServe支持的格式（.mar文件）。

  torch-model-archiver --model-name my_model \
                       --version 1.0 \
                       --serialized-file model_state_dict.pth \
                       --handler image_classifier  # 根据模型类型选择handler（如text_classifier、image_classifier）
  

• 启动TorchServe：

  torchserve --start --model-store ./model_store --models my_model.mar
  

  访问http://localhost:8080/predictions/my_model即可发送预测请求。

关键注意事项

• 依赖一致性：确保训练与部署环境的PyTorch版本、CUDA版本一致，避免兼容性问题。
• 模型格式选择：状态字典更灵活，适合需要修改模型结构的场景；完整模型适合快速部署。
• 性能优化：使用torch.jit.script将模型转换为TorchScript格式，进一步提升推理速度（尤其适合CPU环境）。
• 安全防护：生产环境需添加身份认证、请求限流等安全措施，避免API被滥用。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！