load_model加载含KerasLayer模型报错及预测异常排查

• 训练环境：Google Colab，TensorFlow 2.8.2，基于TensorFlow Hub加载ResNet_v2_50预训练特征提取层搭建二分类图像分类模型，训练、在线推理均运行正常，模型构建代码如下：

import tensorflow as tf
import tensorflow_hub as hub # Provides pretrained models

resnet_url = "https://tfhub.dev/google/imagenet/resnet_v2_50/feature_vector/5"
# Download ResNet model and save as Keras layer
# Trainable : False means we don't want to train it further
resnet_layer = hub.KerasLayer(resnet_url,
                            trainable=False,
                            input_shape=(256,256,3))

# Create model
resnet_model=tf.keras.Sequential([
    # Puts images through downloaded model first
    resnet_layer,
    # Define we will use 2 classes
    tf.keras.layers.Dense(2,
        activation="softmax")
])

• 部署环境：本地Ubuntu，TensorFlow 2.9.1，使用Flask搭建推理服务，加载训练环境导出的h5格式模型文件，服务核心代码如下：

import base64
import numpy as np
import io
import os
from PIL import Image
# 以下为混用的独立keras导入
import keras
from keras import backend as K
from keras.models import Sequential
from keras.models import load_model
from keras.preprocessing.image import ImageDataGenerator
import tensorflow as tf
import tensorflow_hub as hub

from tensorflow.keras.utils import img_to_array
from flask import request
from flask import jsonify
from flask import Flask

app = Flask(__name__)

def get_model():
    global model
    model = load_model('resnet_model_1.h5')
    print(" * Model loaded!")

def preprocess_image(image, target_size):
    if image.mode != "RGB":
        image = image.convert("RGB")
    image = image.resize(target_size)
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)

    return image

print(" * Loading Keras model...")
get_model()

@app.route("/predict", methods=["POST"])
def predict():
    message = request.get_json(force=True)
    encoded = message['image']
    decoded = base64.b64decode(encoded)
    image = Image.open(io.BytesIO(decoded))
    processed_image = preprocess_image(image, target_size=(256, 256))
    
    prediction = model.predict(processed_image).tolist()

    response = {
        'prediction': {
            'dog': prediction[0][0],
            'cat': prediction[0][1]
        }
    }
    return jsonify(response)

• 遇到的问题：
  1. 初始加载模型直接抛出错误：ValueError: Unknown layer: KerasLayer. Please ensure this object is passed to the custom_objects argument
  2. 参考公开方案在load_model时传入参数custom_objects={'KerasLayer':hub.KerasLayer}后，模型可正常加载，但所有输入图片均被预测为cat类别，预测结果完全异常。

按优先级从高到低排查：

• 导入逻辑一致性：是否混用独立Keras和tensorflow.keras两套接口，导致层命名空间不匹配、权重加载错位
• 预处理逻辑一致性：训练阶段和部署阶段的图像预处理流程是否完全对齐，尤其是像素值归一化操作
• 模型存储兼容性：h5格式对TensorFlow Hub自定义层的权重存储是否存在跨版本丢失问题
• 类别映射正确性：训练时的类别标签顺序和推理时的返回索引是否一一对应
• 版本兼容性：跨小版本TensorFlow、TensorFlow Hub是否存在算子实现差异

• 修正导入逻辑，统一使用tensorflow.keras接口
  训练时使用的是tensorflow.keras API，部署代码中混用独立keras包和tensorflow.keras接口会导致层命名空间不匹配，即使模型加载无报错，也会出现权重映射错位的问题。删除所有独立Keras的导入语句，所有Keras相关接口统一从tensorflow.keras导入：

  # 删除以下独立keras的导入代码
  # import keras
  # from keras import backend as K
  # from keras.models import Sequential
  # from keras.models import load_model
  # from keras.preprocessing.image import ImageDataGenerator
  # 替换为统一的tf.keras导入
  from tensorflow.keras.models import Sequential, load_model
  

• 对齐图像预处理逻辑
  ResNet_v2系列预训练模型要求输入像素值归一化到[0,1]区间，这是全量预测为同一类问题的最常见诱因。当前部署代码仅做了图像resize和维度扩展，缺少和训练阶段一致的归一化操作，需要在预处理中补全：

  def preprocess_image(image, target_size):
      if image.mode != "RGB":
          image = image.convert("RGB")
      image = image.resize(target_size)
      image = img_to_array(image)
      # 补全和训练流程一致的归一化逻辑，将0-255范围的像素值缩放到0-1
      image = image / 255.0
      image = np.expand_dims(image, axis=0)
      return image
  

  注意如果训练时使用了tf.keras.applications.resnet_v2.preprocess_input做预处理，部署时需要直接调用该接口处理输入，保证预处理逻辑完全一致。
• 更换模型存储格式，解决跨版本兼容问题
  h5格式对TensorFlow Hub自定义层的权重存储存在跨版本兼容缺陷，保存模型时不要使用h5格式，优先使用TensorFlow原生SavedModel格式导出：

  # 训练环境保存模型时替换原有h5保存逻辑
  resnet_model.save("resnet_saved_model", save_format="tf")
  

  部署环境加载模型时直接加载SavedModel目录即可，不需要额外传入custom_objects参数，所有预训练层和自定义层的权重都会被完整加载。
• 核对类别映射关系
  如果训练时使用flow_from_directory等接口从文件夹生成数据集，类别索引默认按文件夹名字母序排序，需要核对cat、dog类别对应的输出神经元索引，避免和推理接口返回的映射关系错位。
• 版本对齐兜底
  如果上述操作后结果仍有异常，将本地部署环境的TensorFlow、TensorFlow Hub版本调整为和Colab训练环境完全一致，排除小版本算子实现差异带来的计算错误。

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