PyTorch在Ubuntu上的模型部署

在Ubuntu系統上部署PyTorch模型可以通過多種方式實現，以下是一個基本的步驟指南：

環境準備

1. 安裝Python和pip：

sudo apt update
sudo apt install python3 python3-pip

2. 安裝PyTorch：

• 使用pip安裝PyTorch（推薦使用conda安裝，以便于管理依賴）：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

• 或者使用pip安裝CPU版本的PyTorch：

pip install torch torchvision torchaudio

3. 驗證安裝：

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

模型準備

1. 保存訓練好的模型：

# 假設模型已經訓練完成
torch.save(model.state_dict(), 'model.pth')

部署步驟

1. 使用Flask創建一個簡單的Web服務：

• 安裝Flask：

pip install flask

• 創建一個名為app.py的文件，并添加以下代碼：

from flask import Flask, request, jsonify
import torch
import torchvision.transforms as transforms
from PIL import Image

app = Flask(__name__)

# 加載模型
model = torch.load('model.pth', map_location=torch.device('cpu'))
model.eval()

# 定義圖像預處理
transform = 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():
    file = request.files['image']
    image = Image.open(file.stream)
    image = transform(image).unsqueeze(0)  # 添加batch維度
    with torch.no_grad():
        output = model(image)
    _, predicted_idx = torch.max(output, 1)
    return jsonify({'prediction': int(predicted_idx.item())})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

2. 運行Flask應用：

python app.py

3. 測試模型：

• 使用curl發送POST請求：

curl -X POST -F "image=@path_to_your_image.jpg" http://localhost:5000/predict

• 或者使用Postman進行測試。

4. 部署到生產環境：

• 使用Gunicorn：

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 requirements.txt
    RUN pip install -r requirements.txt
    COPY . .
    CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "app:app"]
    

  • 構建Docker鏡像：

    docker build -t your-model-app .
    

  • 運行Docker容器：

    docker run -p 5000:5000 your-model-app
    

通過以上步驟，你可以在Ubuntu系統下成功部署PyTorch模型。根據你的具體需求，可以選擇使用TorchServe、Flask或Django等進行部署。