为什么带自定义层的TensorFlow保存模型性能与训练时差异较大？

问题根因 #

导致加载后模型性能暴跌的核心原因是你的自定义层AutoSRELU不符合Keras序列化规范，训练好的权重没有被正确保存，加载时被重新随机初始化了，具体问题如下：

• 权重创建位置错误：你将k1、k2两个可训练权重的创建逻辑放在了层的__init__方法中，Keras自定义层的权重必须放在build()方法中实现，才能在模型保存时被正确识别、持久化。
• 缺少get_config方法：自定义层必须实现get_config()方法，返回层的初始化参数，才能保证加载时可以正确还原层的配置。
• MinMaxConstraint约束实现错误：约束类的__call__方法没有返回处理后的权重值，导致约束完全不生效，如果你后续要启用这个约束，必须修复返回逻辑。
• 层实例复用问题：你创建了一个auto_srelu实例，先后两次添加到模型中，如果不是特意要共享两个层的权重，这种写法本身也会导致预期外的权重逻辑。
• 评估标签不一致：训练完成后评估用的标签是y_train_pw，加载模型后评估用的是y_train，需要先确认两个变量的内容完全一致，排除标签错配的问题。

修复方案 #

1. 修复自定义约束实现 #

class MinMaxConstraint(keras.constraints.Constraint):
    def __init__(self, minval, maxval):
        self.minval = tf.constant(minval ,dtype='float32')
        self.maxval = tf.constant(maxval ,dtype='float32')
    def __call__(self, w):
        # 增加return返回约束后的权重，用clip_by_value简化逻辑
        return tf.clip_by_value(w, self.minval, self.maxval)
    def get_config(self):
        # 字段名和__init__参数名保持一致
        return {'minval': self.minval.numpy(), 'maxval':self.maxval.numpy()}

2. 修复AutoSRELU自定义层实现 #

class AutoSRELU(keras.layers.Layer):
    def __init__(self, trainable = True, **kwargs):
        super(AutoSRELU, self).__init__(**kwargs) # 传递kwargs给父类
        self.trainable = trainable
        # 不要在__init__中创建权重

    def build(self, input_shape):
        # 权重创建移到build方法中
        initializer0 = keras.initializers.RandomUniform(minval = -1, maxval =1)
        initializer1 = keras.initializers.RandomUniform(minval = 0.5, maxval =3)
        self.k1 = self.add_weight(name='k', shape = (), initializer=initializer0, trainable=self.trainable)
        self.k2 = self.add_weight(name='n', shape = (), initializer=initializer1, trainable=self.trainable)
        super().build(input_shape)

    def call(self, inputs):
        return srelu(inputs, self.k1, self.k2)

    # 新增get_config方法存储层配置
    def get_config(self):
        config = super().get_config()
        config.update({
            "trainable": self.trainable,
        })
        return config

3. 修正模型构建逻辑 #

如果不需要共享两个激活层的权重，分别创建两个独立实例：

model_scratch = models.Sequential()
model_scratch.add(Dense(80, activation='relu', input_dim=(22)))
model_scratch.add(Dense(100, activation='relu'))
model_scratch.add(Dense(120, activation='relu'))
model_scratch.add(Dense(100, activation='relu'))
# 分别创建两个独立的激活层实例
model_scratch.add(AutoSRELU())
model_scratch.add(Dense(120, activation='relu'))
model_scratch.add(AutoSRELU())
model_scratch.add(BatchNormalization())
model_scratch.add(Dense(15, activation='softmax'))

4. 优先使用SavedModel格式保存 #

TensorFlow官方的SavedModel格式对自定义层的支持更完善，不会出现HDF5格式的序列化丢失问题：

# 保存时不要加.h5后缀，会自动生成存储文件夹
model_scratch.save('model_scratch')
# 加载时直接加载文件夹即可
pw_scratch = tf.keras.models.load_model('model_scratch', custom_objects=dependencies)

5. 确认评估数据一致性 #

确认两次评估用的特征X_train、标签y_train_pw和y_train完全一致，排除数据错配的问题。

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