PyTorch在CentOS上的模型部署方法如下:
安装依赖
sudo yum install python3 python3-pip virtualenv。pip3 install torch torchvision torchaudio;CUDA版:添加--extra-index-url参数)。准备模型
torch.save()保存模型为.pth文件。编写服务代码
from flask import Flask, request, jsonify
import torch
app = Flask(__name__)
model = torch.load('model.pth')
model.eval()
@app.route('/predict', methods=['POST'])
def predict():
data = request.json['data']
input_tensor = torch.tensor(data)
with torch.no_grad():
output = model(input_tensor)
return jsonify(output.tolist())
部署服务
nohup或gunicorn运行服务脚本,确保后台运行。docker run启动容器。优化与安全
参考来源: