在Flask或Node中部署大型ML模型：如何实现模型仅加载一次？

不管是Flask还是Node.js，核心思路都是在应用启动阶段一次性加载模型，将模型实例存储在全局/模块级变量中，后续所有请求直接复用这个已加载的模型，避免每次请求重复加载导致的内存浪费和性能损耗。

Flask 实现方案 #

Flask的WSGI运行机制允许我们在应用初始化阶段完成模型加载，所有后续请求共享同一个模型实例。

示例代码 #

from flask import Flask, request, jsonify
import tensorflow as tf  # 替换为你的ML库，如PyTorch、Scikit-learn等

app = Flask(__name__)

# 全局变量存储模型实例
model = None

def load_model():
    """在应用启动时加载模型"""
    global model
    # 加载大体积模型，示例为TensorFlow模型
    model = tf.keras.models.load_model("/path/to/your/large/model")
    print("模型加载完成")

# 在应用启动时执行模型加载（本地运行或生产环境启动时触发）
load_model()

@app.route('/predict', methods=['POST'])
def predict():
    """处理预测请求"""
    if not model:
        return jsonify({"error": "模型未加载完成"}), 500
    
    input_data = request.get_json().get("input")
    if not input_data:
        return jsonify({"error": "缺少输入数据"}), 400
    
    # 使用已加载的模型执行预测
    prediction = model.predict([input_data])
    return jsonify({"prediction": prediction.tolist()})

if __name__ == '__main__':
    # 本地开发时启动服务，生产环境建议用Gunicorn等WSGI服务器
    app.run(host='0.0.0.0', port=5000, threaded=True)

关键注意事项 #

• 生产环境使用多worker服务器（如Gunicorn）时，每个worker会独立加载一次模型。如果模型体积过大，要控制worker数量，避免内存耗尽；或改用单worker + 多线程模式（通过--threads参数配置）。
• 早期Flask版本的@app.before_first_request装饰器已被标记为废弃，建议直接在应用启动阶段加载模型。

Node.js 实现方案 #

Node.js基于单线程事件循环，在应用启动阶段加载模型后，所有请求都会共享这个模型实例，天然适合单实例复用场景。

示例代码 #

const express = require('express');
const tf = require('@tensorflow/tfjs-node'); // 替换为你的ML库，如onnxruntime-node等
const app = express();

// 全局变量存储模型实例
let model = null;

async function initModel() {
    /**初始化加载模型*/
    try {
        model = await tf.loadLayersModel('file:///path/to/your/large/model/model.json');
        console.log('模型加载完成');
    } catch (err) {
        console.error('模型加载失败:', err);
        process.exit(1);
    }
}

// 先加载模型，再启动服务器
initModel().then(() => {
    app.use(express.json());

    app.post('/predict', async (req, res) => {
        if (!model) {
            return res.status(500).json({ error: '模型未加载完成' });
        }

        const inputData = req.body.input;
        if (!inputData) {
            return res.status(400).json({ error: '缺少输入数据' });
        }

        // 使用已加载的模型执行预测
        const tensor = tf.tensor2d([inputData]);
        const prediction = await model.predict(tensor).data();
        tensor.dispose(); // 释放Tensor内存

        res.json({ prediction: Array.from(prediction) });
    });

    const port = process.env.PORT || 3000;
    app.listen(port, () => {
        console.log(`服务器运行在端口 ${port}`);
    });
});

关键注意事项 #

• 如果使用PM2等进程管理器启动多进程，每个进程会独立加载模型，需根据内存资源限制进程数量。
• 对于支持多线程的ML库（如onnxruntime-node），可结合Worker Threads实现多线程推理，但模型实例仍只需加载一次。

通用优化建议 #

• 模型预热：加载完成后执行一次空预测，触发模型内部的懒加载逻辑，避免第一个请求响应延迟过高。
• 内存监控：部署后监控服务器内存使用，避免模型占用内存超出限制导致OOM（内存不足）。
• 容器化部署：用Docker等容器工具部署时，设置合理的内存配额，确保容器有足够资源加载模型。

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