如何在Flask Web框架中加载多个Pickle模型并在Web API中输出预测结果

核心改造逻辑 #

替换原有单模型全局加载的实现，采用注册表字典结构统一管理所有已加载的Pickle模型，应用启动时批量扫描加载指定目录的模型文件，Web端新增模型选择交互入口，预测阶段根据用户选中的模型完成推理并返回结果。

具体调整点 #

• 用字典结构存储所有模型实例，替代原单个全局model变量，实现多模型的快速索引调用
• 新增批量加载函数，应用启动时自动扫描指定模型目录下的所有.pkl文件，无需手动逐行编写加载代码
• 适配表单结构变化：过滤非数值类型的表单输入，避免把模型选择字段的字符串值误判为预测特征，触发类型转换错误
• 新增模型合法性校验逻辑，避免用户传入无效模型名称时触发程序崩溃

改造后的完整可运行代码 #

import os
import numpy as np
from flask import Flask, request, render_template
import pickle

app = Flask(__name__)
# 模型注册表：键为模型名称，值为加载完成的模型实例
MODEL_REGISTRY = {}
# 配置存放Pickle模型的目录路径
MODEL_DIR = "./models"

def load_all_models():
    """应用启动时批量加载目录下所有Pickle模型"""
    for file_name in os.listdir(MODEL_DIR):
        if file_name.endswith(".pkl"):
            model_path = os.path.join(MODEL_DIR, file_name)
            # 去掉文件后缀作为模型的唯一标识
            model_key = os.path.splitext(file_name)[0]
            with open(model_path, "rb") as f:
                MODEL_REGISTRY[model_key] = pickle.load(f)
    print(f"加载完成，当前可用模型列表：{list(MODEL_REGISTRY.keys())}")

# 应用启动前完成所有模型加载
load_all_models()

@app.route('/')
def home():
    # 将可用模型列表传递给前端，供下拉选择使用
    return render_template('index.html', model_list=list(MODEL_REGISTRY.keys()))

@app.route('/predict', methods=['POST'])
def predict():
    # 获取用户选中的模型名称
    selected_model = request.form.get("selected_model")
    # 校验所选模型是否存在
    if selected_model not in MODEL_REGISTRY:
        return render_template('index.html', prediction_text="所选模型不存在，请刷新重试")
    
    # 提取数值型输入作为预测特征，过滤非数值的模型选择字段
    int_features = [int(x) for x in request.form.values() if x.isdigit()]
    final_features = [np.array(int_features)]
    # 从注册表取出对应模型完成预测
    model = MODEL_REGISTRY[selected_model]
    prediction = model.predict(final_features)
    output = round(prediction[0], 2)

    return render_template(
        'index.html',
        prediction_text=f'模型[{selected_model}]的水量预测结果为{output}升'
    )

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

前端页面配套调整 #

在原有index.html的预测表单区域新增模型选择模块，核心代码片段：

<div>
    <label>选择预测模型：</label>
    <select name="selected_model" required>
        {% for model in model_list %}
        <option value="{{ model }}">{{ model }}</option>
        {% endfor %}
    </select>
</div>

部署注意事项 #

• 在项目根目录新建models文件夹，将所有待部署的Pickle模型统一放入该目录，应用启动时会自动遍历加载
• 原单模型实现存在缩进缺失问题（predict函数内部逻辑、app.run代码未缩进），直接运行会触发语法错误，改造版本已同步修正
• 模型加载完成后可查看控制台输出的可用模型列表，确认所有模型是否加载成功

内容的提问来源于stack exchange，提问作者Shridatha Hegde