TensorFlow模型生产推理：如何避免重复加载以提升效率？

你的脚本失效的常见原因 #

你写的全局变量缓存脚本在本地单进程测试可能有效，但放到生产环境（尤其是多进程Web服务）会失效，核心原因是：

• 多进程环境下，每个进程拥有独立的内存空间，全局变量无法跨进程共享，每个worker进程都会重新加载一次模型。
• 如果是CGI这类每次请求启动新进程的环境，全局变量会被反复初始化，根本无法缓存。

核心优化策略与实践 #

1. 单进程服务：利用启动事件加载模型 #

在Flask、FastAPI这类单进程Web框架中，通过服务启动钩子加载模型，确保模型只初始化一次，所有请求复用同一个实例。

以FastAPI为例：

import tensorflow as tf
from fastapi import FastAPI
from pydantic import BaseModel

# 定义输入数据结构
class InferenceRequest(BaseModel):
    input_data: list[list[float]]

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

app = FastAPI()

# 服务启动时加载模型，仅执行一次
@app.on_event("startup")
def load_model():
    global model
    # 推荐使用SavedModel格式，比.h5更适合生产
    model = tf.keras.models.load_model("saved_model/my_model")

@app.post("/predict")
def predict(request: InferenceRequest):
    input_tensor = tf.convert_to_tensor(request.input_data, dtype=tf.float32)
    predictions = model.predict(input_tensor)
    return {"predictions": predictions.tolist()}

Flask的实现类似，使用启动逻辑加载模型：

from flask import Flask, request, jsonify
import tensorflow as tf

app = Flask(__name__)
model = None

@app.before_first_request
def load_model():
    global model
    model = tf.keras.models.load_model("saved_model/my_model")

@app.route("/predict", methods=["POST"])
def predict():
    input_data = request.json["input_data"]
    input_tensor = tf.convert_to_tensor(input_data)
    predictions = model.predict(input_tensor)
    return jsonify({"predictions": predictions.tolist()})

2. 多进程/高并发场景：使用专业模型服务框架 #

生产环境高并发场景下，不要自己实现缓存，优先用TensorFlow官方的TensorFlow Serving，或者通用的Triton Inference Server，这些框架原生解决模型加载、缓存、版本管理和高并发问题。

步骤1：导出模型为SavedModel格式

先把你的.h5模型转换为生产友好的SavedModel：

model = tf.keras.models.load_model("path_to_model.h5")
# 导出到指定目录，目录名`1`代表模型版本，方便后续热更新
model.save("saved_model/my_model/1")

步骤2：启动TensorFlow Serving服务

# 启动服务，指定端口和模型路径
tensorflow_model_server --port=8501 --model_name=my_model --model_base_path=/path/to/saved_model/my_model

步骤3：客户端请求示例

import requests
import json

input_data = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
payload = json.dumps({"instances": input_data})
response = requests.post(
    "http://localhost:8501/v1/models/my_model:predict",
    data=payload,
    headers={"Content-Type": "application/json"}
)
predictions = json.loads(response.text)["predictions"]

3. 进程内缓存：单例模式实现（仅单进程场景） #

如果必须在单进程脚本中实现缓存，可以用单例类确保模型只被实例化一次：

import tensorflow as tf

class CachedModel:
    _instance = None

    @classmethod
    def get_model(cls):
        if cls._instance is None:
            cls._instance = tf.keras.models.load_model("saved_model/my_model")
        return cls._instance

# 推理函数直接复用单例模型
def inference(input_data):
    model = CachedModel.get_model()
    return model.predict(input_data)

4. 进阶优化：结合推理加速 #

加载模型时同时完成推理优化，进一步提升性能：

• 转换为TensorRT引擎（适合NVIDIA GPU加速）：

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

# 转换模型为TensorRT优化版本
converter = trt.TrtGraphConverterV2(input_saved_model_dir="saved_model/my_model")
converter.convert()
converter.save("saved_model/my_model_trt")

# 加载优化后的模型，后续复用
model = tf.keras.models.load_model("saved_model/my_model_trt")

关键最佳实践总结 #

• 优先用专业模型服务框架：TensorFlow Serving、Triton Inference Server是生产环境的首选，无需自己维护缓存逻辑。
• 避免在请求时加载模型：所有模型加载逻辑必须放在服务启动阶段，而非请求处理流程中。
• 使用SavedModel格式：相比.h5，SavedModel支持更多生产特性，兼容性更好。
• 多进程环境禁用进程内缓存：多进程下进程内缓存会导致每个进程重复加载模型，反而浪费资源，改用集中式模型服务。

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