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被滥用。

0 赞

0 踩