Flask 推理 API：把模型变成服务

模型训练完了，总不能永远躺在 .pt 或 .bin 文件里吧？你得给它搭个“桥”，让外界能调用它。这个桥就是推理 API。

我习惯用 Flask。它轻量、灵活，特别适合快速把模型包装成服务。你想想看，生产环境里可能用 FastAPI 或 gRPC，但原型验证、小规模部署、内部工具——Flask 完全够用，而且调试起来特别顺手。

本章核心：用最少代码，把推理逻辑变成 HTTP 接口。

4.1 路由设计：API 的“门牌号”

每一个 API 接口就是一个 URL 路径。我通常把推理接口命名为 /predict 或 /v1/predict，带上版本号方便迭代。

from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    # 推理逻辑
    return jsonify({'result': 'ok'})

这里我用了 POST 方法。为什么呢？推理请求通常带输入数据，GET 方法有长度限制，而且数据会暴露在 URL 里，不合适。

个人习惯：我会在路由里加版本号，比如 /v1/predict。万一以后模型升级、输入输出格式变了，新旧版本可以共存，客户端平滑迁移。

4.2 请求解析：从 HTTP 里“拧出”数据

客户端发来请求，数据通常在请求体里。最常见的是 JSON 格式。你得把它解析成 Python 对象。

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json()
    if not data:
        return jsonify({'error': '请求体不能为空'}), 400
    
    input_text = data.get('text')
    if not input_text:
        return jsonify({'error': '缺少 text 字段'}), 400
    
    # 后续推理...

这里要注意：request.get_json() 可能返回 None。我在项目中遇到过客户端忘记设 Content-Type: application/json，结果 Flask 解析不了，返回了空。所以校验不能省。

有时候输入是图片或音频，传的是 base64 编码的字符串，或者直接上传文件。Flask 处理上传文件也方便：

if 'file' not in request.files:
    return jsonify({'error': '缺少文件'}), 400

file = request.files['file']
image_bytes = file.read()
# 后续处理...

我曾经踩过的坑：文件上传没设最大大小限制，结果有人传了个 2GB 的视频，直接把内存撑爆了。记得配置 app.config['MAX_CONTENT_LENGTH']。

4.3 模型加载：全局单例与预热

模型加载是高频操作。你不能每次请求都重新加载一次，那太慢了。我通常把模型做成全局变量，在服务启动时加载一次。

model = None

def load_model():
    global model
    print("加载模型中...")
    model = torch.load('model.pt', map_location='cpu')
    model.eval()
    print("模型加载完成")

@app.before_first_request
def warm_up():
    load_model()

@before_first_request 会在第一个请求到来前执行。如果模型很大，第一次请求会有点慢。我后来改用 if __name__ == '__main__' 里直接加载，让服务在启动时就准备好。

if __name__ == '__main__':
    load_model()   # 启动时加载
    app.run(host='0.0.0.0', port=5000)

多 GPU 场景呢？模型可以加载到指定设备：

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)

但注意，Flask 默认是同步的。如果模型在 GPU 上推理，一个请求没处理完，下一个会排队。高并发时建议配合 Gunicorn 多 worker 或异步框架。

关键点：模型全局加载一次，避免重复 I/O。GPU 显存有限，留意 worker 数量和显存占用。

4.4 响应封装：统一格式，方便调用方

我见过很多 API，输入格式五花八门，输出格式各自为政。调用方维护起来特别痛苦。我个人习惯统一响应结构：

{
    "code": 0,
    "message": "success",
    "data": {
        "prediction": "cat",
        "confidence": 0.98
    }
}

错误时：

{
    "code": 1001,
    "message": "模型推理失败",
    "data": null
}

封装成一个函数：

def success_response(data, message="success"):
    return jsonify({
        'code': 0,
        'message': message,
        'data': data
    })

def error_response(code, message):
    return jsonify({
        'code': code,
        'message': message,
        'data': None
    }), 200  # 或者 400/500

这样客户端解析就简单了：先看 code，0 代表成功，其他数字查错误表。

避坑：HTTP 状态码我通常统一返回 200，错误信息放在业务 code 里。为什么？因为有些代理或 CDN 会缓存非 200 的响应，导致客户端拿到旧数据。不过这只是我的做法，你们按团队规范来。

4.5 错误处理：让失败也优雅

推理时可能遇到各种问题：输入格式不对、模型前向传播报错、显存不足……你不希望服务直接返回 500 还带一堆堆栈信息吧？

Flask 提供了错误处理器：

@app.errorhandler(400)
def bad_request(e):
    return error_response(400, "请求参数错误")

@app.errorhandler(404)
def not_found(e):
    return error_response(404, "接口不存在")

@app.errorhandler(500)
def server_error(e):
    print(f"内部错误: {e}")  # 本地打印日志
    return error_response(500, "服务器内部错误")

更常见的是在推理逻辑里用 try-except：

try:
    with torch.no_grad():
        output = model(input_tensor)
except RuntimeError as e:
    if 'out of memory' in str(e):
        torch.cuda.empty_cache()
        return error_response(1002, "显存不足，请稍后重试")
    return error_response(1001, f"推理失败: {str(e)}")

你看，错误处理不是简单返回异常，而是给出具体原因。调用方可以根据 code 做重试或提示用户。

我曾翻过的车：显存 OOM 后没清缓存，后续请求依然失败。所以 torch.cuda.empty_cache() 一定要调，虽然不保证立竿见影，但至少试试。

4.6 完整代码骨架与流程

把上面所有东西拼起来，一个可用的推理 API 就成型了：

import torch
import torch.nn.functional as F
from flask import Flask, request, jsonify

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB

model = None
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

def load_model():
    global model
    print("Loading model...")
    model = torch.load('model.pt', map_location=device)
    model.eval()
    print("Model loaded on", device)

def success_response(data, message="success"):
    return jsonify({'code': 0, 'message': message, 'data': data})

def error_response(code, message):
    return jsonify({'code': code, 'message': message, 'data': None})

@app.route('/v1/predict', methods=['POST'])
def predict():
    data = request.get_json()
    if not data or 'text' not in data:
        return error_response(400, "缺少 text 字段")
    
    input_text = data['text']
    # 假设模型接收 tokenized 输入
    tokens = tokenizer.encode(input_text, return_tensors='pt').to(device)
    
    try:
        with torch.no_grad():
            logits = model(tokens)
            pred = torch.argmax(F.softmax(logits, dim=-1), dim=-1).item()
        return success_response({'prediction': pred, 'confidence': 0.95})
    except Exception as e:
        print(f"Predict error: {e}")
        return error_response(1001, "推理失败")

@app.errorhandler(400)
def handle_400(e):
    return error_response(400, "请求参数错误")

@app.errorhandler(500)
def handle_500(e):
    return error_response(500, "服务器内部错误")

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

你想想看，是不是很清晰？路由接收请求，解析数据，调用模型，封装响应，处理异常——一条龙。

4.7 请求处理流程全景

下面这张图概括了整个请求的生命周期，从客户端发起到返回结果。

Flask 推理 API 请求处理流程 客户端请求 路由匹配 请求解析 get_json() / files 模型推理 响应封装 错误处理 响应 异常 说明：模型为全局单例， 启动时加载，避免重复 I/O。

看到那个异常分支了吗？模型推理阶段如果抛出异常，不会直接挂掉，而是流入错误处理节点，统一返回结构化错误信息。这个设计让服务更健壮。

4.8 本章小结

Flask 搭建推理 API 其实就这五步：路由、解析、加载、响应、错误处理。说白了，就是让外部系统能用 HTTP 跟你的模型聊天。

我个人觉得最难的部分不是写代码，而是想清楚边界情况：输入为空怎么办？模型还没加载完怎么办？CPU 和 GPU 怎么切换？把这些都想明白，API 就稳了。

记住：一个推理 API 的价值不在于用了多先进的模型，而在于调用方能不能无脑接入。统一格式、认真校验、优雅报错——这就是专业。

← 上一章 📖 返回目录 下一章 →