PyTorch在Ubuntu上的模型部署

在Ubuntu系统上部署PyTorch模型可以通过多种方式实现，以下是一个基本的步骤指南：

环境准备

安装Python和pip：

sudo apt update
sudo apt install python3 python3-pip

安装PyTorch：

使用pip安装PyTorch（推荐使用conda安装，以便于管理依赖）：

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

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

pip install torch torchvision torchaudio

验证安装：

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

模型准备

保存训练好的模型：

# 假设模型已经训练完成
torch.save(model.state_dict(), 'model.pth')

部署步骤

使用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)

运行Flask应用：

python app.py

测试模型：

使用curl发送POST请求：

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

或者使用Postman进行测试。

部署到生产环境：

使用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等进行部署。

环境准备

模型准备

部署步骤

最新问答

相关标签