基于Flask的模型推理服务部署指南

基于Flask的模型推理服务部署指南

最近在尝试将训练好的大模型部署为推理服务时，踩了不少坑，特此记录下完整的部署流程和关键注意事项。

环境准备

首先确保Python环境已安装必要的依赖：

pip install flask torch torchvision

核心代码实现

创建一个简单的Flask应用来提供推理服务：

from flask import Flask, request, jsonify
import torch
import json

app = Flask(__name__)
model = None

# 加载模型
@app.route('/load_model', methods=['POST'])
def load_model():
    global model
    # 这里需要根据实际模型路径修改
    model = torch.load('path/to/your/model.pth')
    model.eval()
    return jsonify({'status': 'model loaded'})

# 推理接口
@app.route('/predict', methods=['POST'])
def predict():
    global model
    if model is None:
        return jsonify({'error': 'model not loaded'}), 400
    
    data = request.get_json()
    input_tensor = torch.tensor(data['input'])
    with torch.no_grad():
        output = model(input_tensor)
    return jsonify({'output': output.tolist()})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

部署注意事项

1. 模型加载优化：避免每次请求都重新加载模型，建议使用单例模式。
2. 内存管理：对于大模型，注意设置torch.cuda.empty_cache()防止显存溢出。
3. 并发处理：生产环境应配合Gunicorn或uWSGI进行部署，而非直接用Flask内置服务器。

本地测试

启动服务后使用curl测试：

curl -X POST http://localhost:5000/predict \
     -H "Content-Type: application/json" \
     -d '{"input": [1,2,3]}'

希望这份指南能帮助到同样在做模型部署的朋友们！