ubuntu

Ubuntu下PyTorch模型怎么部署

小樊
48
2025-05-10 04:32:48
栏目: 智能运维

在Ubuntu系统下部署PyTorch模型,可以遵循以下步骤:

1. 安装PyTorch

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

使用pip安装:

pip install torch torchvision

使用conda安装:

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch

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:

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

6. 部署到生产环境

如果你需要将模型部署到生产环境,可以考虑使用以下方法:

使用Gunicorn:

安装Gunicorn:

pip install gunicorn

运行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模型。

0
看了该问题的人还看了