Keras 2转Keras 3后模型加载及predict方法调用问题求助

问题根源 #

Keras 3是多后端统一框架，你去年用Keras 2（tf.keras）保存的模型属于TensorFlow SavedModel格式，直接用tf.keras.models.load_model加载会触发兼容性报错，而keras.layers.TFSMLayer是Keras 3提供的适配TensorFlow SavedModel的层组件——注意它是层不是完整模型，因此没有predict方法。

方案1：直接调用TFSMLayer实例（最快修复） #

修改预测逻辑，直接调用TFSMLayer对象（层的调用方式），而不是用predict。TFSMLayer的返回值是字典，需要根据模型的输出节点名称提取结果：

import keras
import numpy as np

model = keras.layers.TFSMLayer('UNetW_bs32_ep25_s256_aug1_model', call_endpoint='serving_default')
image_size = 256

def predict_image(image, model=model):
    # Preprocess image
    image = image / 255.
    image = np.expand_dims(image, axis=0)
    # 直接调用层，获取输出字典
    output_dict = model(image)
    # 替换为你模型的输出节点名称（查看方法见下文）
    mask = output_dict['your_output_node_name'][0]
    mask = np.round(mask)
    # Postprocess mask
    mask = np.squeeze(mask, axis=-1)
    return mask

mask = predict_image(image)

查看输出节点名称的方法：打开终端运行以下命令，在SignatureDefs下的serving_default模块里找到outputs对应的键名：

saved_model_cli show --dir UNetW_bs32_ep25_s256_aug1_model --all

方案2：将TFSMLayer包装为Keras模型（恢复predict方法） #

如果习惯用predict方法，可以把TFSMLayer封装成完整的Keras模型：

import keras
import numpy as np
from keras import Input

# 加载TFSMLayer
tf_layer = keras.layers.TFSMLayer('UNetW_bs32_ep25_s256_aug1_model', call_endpoint='serving_default')

# 根据你的输入形状创建输入层（示例为(256,256,3)，替换为实际输入维度）
inputs = Input(shape=(256, 256, 3))
# 通过TFSMLayer传递输入，替换为实际输出节点名
outputs = tf_layer(inputs)['your_output_node_name']
# 构建完整模型
model = keras.Model(inputs=inputs, outputs=outputs)

image_size = 256

def predict_image(image, model=model):
    # Preprocess image
    image = image / 255.
    image = np.expand_dims(image, axis=0)
    # 现在可以正常使用predict
    mask = model.predict(image)[0]
    mask = np.round(mask)
    # Postprocess mask
    mask = np.squeeze(mask, axis=-1)
    return mask

mask = predict_image(image)

方案3：直接用TensorFlow加载模型 #

如果不需要Keras 3的多后端特性，也可以绕过Keras层，直接用TensorFlow加载SavedModel：

import tensorflow as tf
import numpy as np

# 用TensorFlow加载模型
loaded_model = tf.saved_model.load('UNetW_bs32_ep25_s256_aug1_model')
# 获取预测签名
predict_fn = loaded_model.signatures['serving_default']

image_size = 256

def predict_image(image, predict_fn=predict_fn):
    # Preprocess image
    image = image / 255.
    image = np.expand_dims(image, axis=0)
    # 转换为TensorFlow张量
    tf_image = tf.convert_to_tensor(image, dtype=tf.float32)
    # 执行预测
    output_dict = predict_fn(tf_image)
    mask = output_dict['your_output_node_name'].numpy()[0]
    mask = np.round(mask)
    # Postprocess mask
    mask = np.squeeze(mask, axis=-1)
    return mask

mask = predict_image(image)

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