如何降低tensorflow.keras.models.load_model加载.h5模型时的CPU使用率

问题根因 #

导入TensorFlow模块时默认会占用所有可用CPU核心做初始化运算，你的Pod CPU限额仅1核，初始化时CPU直接打满被K8s强制驱逐，导致Pod崩溃。

方案1：保留.h5格式，优化加载逻辑 #

第一步：提前限制TensorFlow线程数 #

在导入任何TensorFlow相关模块前，添加以下环境变量配置：

import os
# 严格限制TensorFlow CPU线程数，匹配Pod 1核的限额
os.environ["TF_NUM_INTEROP_THREADS"] = "1"
os.environ["TF_NUM_INTRAOP_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
# 关闭不必要的日志输出减少资源消耗
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

第二步：加载模型时关闭编译选项 #

推理场景不需要训练用的损失函数、优化器配置，加载时添加compile=False参数大幅降低资源消耗：

from tensorflow.keras.models import load_model
# 仅加载推理需要的权重和计算图
model = load_model("你的模型路径.h5", compile=False)

方案2：转换为TFLite格式（推荐，资源占用降低80%以上） #

TFLite是TensorFlow专门面向推理场景推出的轻量部署格式，不需要安装全量TensorFlow依赖，CPU和内存占用远低于原生.h5模型加载。

第一步：本地转换模型格式 #

import tensorflow as tf
# 加载原h5模型
model = tf.keras.models.load_model("你的模型路径.h5", compile=False)
# 转换为TFLite格式，开启量化进一步压缩体积
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()
# 保存转换后的模型
with open("model.tflite", "wb") as f:
    f.write(tflite_model)

第二步：部署时使用轻量TFLite运行时 #

部署时不需要安装全量tensorflow包，只需要安装tflite-runtime依赖，加载推理代码如下：

import tflite_runtime.interpreter as tflite
import numpy as np

# 服务启动时加载一次模型即可
interpreter = tflite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 推理接口调用
def predict(input_data):
    # 输入数据需要匹配模型要求的维度和数据类型
    interpreter.set_tensor(input_details[0]['index'], input_data.astype(np.float32))
    interpreter.invoke()
    return interpreter.get_tensor(output_details[0]['index'])

额外优化建议 #

• 模型加载逻辑必须放在FastAPI服务启动阶段执行，不要放在接口路由函数中，避免每次请求重复加载模型消耗资源
• 若使用全量TensorFlow，安装时选择tensorflow-cpu包，不要安装带GPU支持的版本，减少依赖体积和初始化资源消耗
• 若仍有偶发CPU峰值导致Pod崩溃，可以将AKS Pod的CPU Request调整为500m，避免调度到资源紧张的节点

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