Ubuntu如何部署PyTorch模型 - 问答

在Ubuntu系统下部署PyTorch模型可以通过以下步骤完成：

1. 安装PyTorch

首先，确保你已经安装了PyTorch。你可以使用pip或conda来安装。

使用pip安装：

pip install torch torchvision torchaudio

使用conda安装（推荐，特别是有NVIDIA GPU时）：
```
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
```
请根据你的CUDA版本选择合适的cudatoolkit版本。

2. 准备模型

确保你的模型已经训练完成并保存为.pt或.pth文件。

3. 创建一个简单的Flask应用

Flask是一个轻量级的Web框架，适合用来部署模型。

安装Flask：
```
pip install flask
```

创建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('path_to_your_model.pt', 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():
    if request.method == 'POST':
        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)

4. 运行Flask应用

在终端中运行以下命令来启动Flask应用：

python app.py

5. 测试模型

你可以使用curl或Postman来测试你的模型。

使用curl：

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

使用Postman：
- 打开Postman。
- 创建一个新的POST请求，URL为http://localhost:5000/predict。
- 在Body中选择form-data，添加一个键为image的文件字段，并上传你的图像文件。
- 发送请求并查看响应。

6. 部署到生产环境（可选）

如果你需要将模型部署到生产环境，可以考虑以下方法：

使用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 run -p 5000:5000 your-model-app

通过以上步骤，你可以在Ubuntu系统下成功部署PyTorch模型。

0 赞

0 踩