TensorFlow 2.6.2加载含自定义层的ViT模型时出现TypeError错误求助

问题原因 #

你的自定义DataAugmentation层的__init__方法未接收Keras反序列化过程中自动传递的name参数，也没有通过**kwargs处理额外参数并传递给父类构造函数，导致模型加载时触发参数不匹配错误。

修改后的自定义层代码 #

import tensorflow as tf
from tensorflow.keras.layers import Layer, Resizing, RandomFlip, RandomRotation, RandomZoom

@tf.keras.utils.register_keras_serializable()
class DataAugmentation(Layer):
    def __init__(self, norm, SIZE, **kwargs):
        super(DataAugmentation, self).__init__(**kwargs)
        self.norm = norm
        self.SIZE = SIZE
        self.resize = Resizing(SIZE, SIZE)
        self.flip = RandomFlip('horizontal')
        self.rotation = RandomRotation(factor=0.02)
        self.zoom = RandomZoom(height_factor=0.2, width_factor=0.2)
    
    def call(self, X):
        x = self.norm(X)
        x = self.resize(x)
        x = self.flip(x)
        x = self.rotation(x)
        x = self.zoom(x)
        return x

    def get_config(self):
        config = super().get_config()
        config.update({
            "norm": self.norm,
            "SIZE": self.SIZE,
        })
        return config

加载模型的正确方式 #

修改层代码后，重新训练并保存模型，加载时可使用以下方式：

• 直接加载（已注册序列化的情况下）：

model = tf.keras.models.load_model("ViT-Model-new.h5")

• 或结合custom_object_scope：

with tf.keras.utils.custom_object_scope({"DataAugmentation": DataAugmentation}):
    model = tf.keras.models.load_model("ViT-Model-new.h5")

为什么之前的方法无效 #

你尝试的注册装饰器和custom_object_scope仅用于让Keras识别自定义层类，但层本身的构造函数不支持接收name等额外参数，因此即使找到类，初始化时仍会触发参数错误。

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