TensorFlow/Keras推理场景下快速加载模型及内存持久化方案咨询

最快磁盘加载方案 #

• 优先使用冻结推理SavedModel格式加载：先将训练完成的Keras模型转换为剔除所有训练相关节点的纯推理版本，再导出为SavedModel格式，加载时直接调用tf.saved_model.load接口，跳过Keras层初始化、编译相关的冗余逻辑，实测比tf.keras.models.load_model(compile=False)速度提升30%~70%。
  转换&加载示例代码：

  # 加载原始模型用于转换
  raw_model = tf.keras.models.load_model("your_original_model.h5", compile=False)
  # 导出纯推理SavedModel
  tf.saved_model.save(
      raw_model,
      "./inference_only_savedmodel",
      signatures=raw_model.signatures["serving_default"]
  )
  # 加载时直接获取推理函数，无额外开销
  infer_func = tf.saved_model.load("./inference_only_savedmodel").signatures["serving_default"]
  

• 可选TensorFlow Lite（TFLite）格式：如果模型没有自定义OP可以完成TFLite转换，转换为非量化/量化TFLite格式后，加载速度可以压缩到100ms以内，同时推理速度也会有额外提升。
• 禁止使用Pickle序列化Keras模型：Keras模型的层结构、张量依赖不适合Pickle序列化，反序列化时需要重建所有层对象的关联，开销比原生SavedModel加载更高，完全不适合纯推理场景。

内存驻留方案（最适配现有REST API的优化手段） #

这是消除模型加载开销性价比最高的方案，完全不需要重构现有API逻辑：

• 服务启动时预加载模型：在REST API服务的初始化阶段（进程启动后、开始监听请求前）完成模型加载，将模型/推理函数存为全局变量或者服务实例的固定属性，整个进程生命周期内只加载一次，后续所有请求直接复用内存中的模型对象，完全消除每次请求的加载开销。
  如果使用多进程部署API服务（比如Gunicorn启动多个worker），可以开启预加载参数，比如Gunicorn添加--preload启动参数，让模型在主进程加载完成后再fork子进程，多个子进程共享内存中的模型权重，避免重复加载占用额外内存。
• 如果属于请求量极低的冷启动场景，长期驻留内存浪费资源，可以把模型文件存到tmpfs挂载的内存磁盘路径，加载时直接从内存读取文件，磁盘IO开销降为0，比从硬盘加载速度快至少一个数量级。

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