基于Flask构建模型推理服务接口

基于Flask构建模型推理服务接口

在大模型推理场景中，将训练好的模型封装成可访问的Web服务是常见的实践。本文将介绍如何使用Flask快速搭建一个模型推理服务接口。

环境准备

pip install flask torch torchvision

核心代码实现

from flask import Flask, request, jsonify
import torch
import torchvision.transforms as transforms
from PIL import Image

app = Flask(__name__)
model = None  # 假设模型已加载

@app.route('/predict', methods=['POST'])
def predict():
    try:
        # 获取图片数据
        file = request.files['image']
        image = Image.open(file.stream).convert('RGB')
        
        # 图像预处理
        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])
        ])
        
        image_tensor = transform(image).unsqueeze(0)
        
        # 模型推理
        with torch.no_grad():
            output = model(image_tensor)
            _, predicted = torch.max(output.data, 1)
            
        return jsonify({'prediction': int(predicted)})
    except Exception as e:
        return jsonify({'error': str(e)}), 500

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

使用方法

1. 启动服务：python app.py
2. 发送POST请求：curl -X POST -F "image=@test.jpg" http://localhost:5000/predict

该方案便于快速验证模型推理能力，适合在开发阶段使用。