在Ubuntu系统下部署PyTorch模型,可以遵循以下步骤:
首先,确保你已经安装了PyTorch。你可以使用pip或conda来安装。
pip install torch torchvision
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
确保你的模型已经训练完成并保存为.pt
或.pth
文件。
Flask是一个轻量级的Web框架,适合用来部署模型。
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('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)
在终端中运行以下命令来启动Flask应用:
python app.py
你可以使用curl
或Postman来测试你的模型。
curl -X POST -F "image=@path_to_your_image.jpg" http://localhost:5000/predict
http://localhost:5000/predict
。form-data
,添加一个键为image
的文件字段,并上传你的图像文件。如果你需要将模型部署到生产环境,可以考虑使用以下方法:
安装Gunicorn:
pip install gunicorn
运行Gunicorn:
gunicorn -w 4 -b 0.0.0.0:5000 app:app
创建一个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模型。