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