保存HDF5模型时遇ValueError: Unknown layer: 'TFOpLambda'问题求助

保存模型为HDF5格式后，加载时触发如下错误：

ValueError: Unknown layer: 'TFOpLambda'. Please ensure you are using a `keras.utils.custom_object_scope` and that this object is included in the scope.

环境版本：

• TensorFlow 2.16.1
• Keras 3.3.3
  加载模型的代码片段：

def predict(image):
    # Load the model using custom object scope
    with tf.keras.utils.custom_object_scope(custom_object_scope()):
        loaded_model = tf.keras.models.load_model('model.h5')

    img_array = np.array(image)
    img = tf.image.resize(img_array, size=(256, 256))
    img = tf.expand_dims(img, axis=0)
    img = img / 255.
    intensity = loaded_model.predict(img)
    return intensity

解决方案 #

1. 切换为SavedModel格式（推荐） #

HDF5对TensorFlow原生操作（如TFOpLambda）的兼容性有限，官方更推荐使用SavedModel格式：

• 保存模型时替换为：

  model.save('my_model')  # 保存为文件夹形式的SavedModel
  

• 加载时无需额外配置，直接调用：

  loaded_model = tf.keras.models.load_model('my_model')
  

2. 显式注册TFOpLambda到自定义对象作用域 #

若必须使用HDF5格式，需将TFOpLambda注册到自定义对象字典中：
首先定义包含目标层的字典：

from tensorflow.keras.layers import TFOpLambda

def get_custom_objects():
    return {'TFOpLambda': TFOpLambda}

然后修改加载代码（替换原未定义的custom_object_scope()）：

def predict(image):
    with tf.keras.utils.custom_object_scope(get_custom_objects()):
        loaded_model = tf.keras.models.load_model('model.h5')
    
    # 后续图像处理与预测逻辑保持不变
    img_array = np.array(image)
    img = tf.image.resize(img_array, size=(256, 256))
    img = tf.expand_dims(img, axis=0)
    img = img / 255.
    intensity = loaded_model.predict(img)
    return intensity

3. 封装TensorFlow原生操作为自定义Keras层 #

如果模型中直接使用了tf.image.resize这类原生操作作为层的一部分，建议封装为自定义层，避免生成TFOpLambda：

class ResizeLayer(tf.keras.layers.Layer):
    def __init__(self, size=(256,256)):
        super().__init__()
        self.size = size
    
    def call(self, inputs):
        return tf.image.resize(inputs, size=self.size)

构建模型时用该自定义层替代直接调用tf.image.resize，即可避免HDF5加载时的未知层问题。

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