自定义Keras层训练与推理输出不一致问题求助

核心问题分析 #

训练时输出固定为0.0或1.0，本质是tanh激活函数进入饱和区，导致梯度消失，权重无法更新，最终输出被卡在极限值。同时自定义层的实现存在几个细节问题，加剧了这个现象：

• 在call方法中动态创建Flatten层，结合@tf.function装饰器，会导致张量形状追踪异常，干扰正常计算流程。
• 强制设置输入形状为(1, self.input_shape_accu)，限制了输入的batch维度，当训练时输入无batch维度时，会引发形状不兼容问题，间接促使模型向极端值收敛。
• 训练循环中输入未显式添加batch维度，与层的预期输入格式不匹配。

修复步骤 #

1. 重构自定义BoundingBoxNumber层 #

修改层的实现，将Flatten层移到__init__中初始化，移除强制形状设置，同时可以考虑用sigmoid直接替代tanh+缩放，简化0-1映射逻辑：

class BoundingBoxNumber(tf.keras.layers.Layer):
    def __init__(self, self_input_shape):
        super(BoundingBoxNumber, self).__init__()
        self.input_shape_custom = self_input_shape
        self.input_shape_accu = 1
        for item in self_input_shape:
            self.input_shape_accu *= item
        
        # 提前初始化Flatten层，避免tf.function内动态创建
        self.flatten_layer = tf.keras.layers.Flatten()
        # 用sigmoid直接输出0-1区间，替代tanh+缩放
        self.internal_dense_layer = tf.keras.layers.Dense(1, activation='sigmoid')

    @tf.function
    def call(self, inputs):
        inputs = self.flatten_layer(inputs)
        # 移除强制形状设置，让TensorFlow自动处理维度
        output = self.internal_dense_layer(inputs)
        return output

如果坚持使用tanh，保留缩放步骤，但必须移除形状强制设置：

class BoundingBoxNumber(tf.keras.layers.Layer):
    def __init__(self, self_input_shape):
        super(BoundingBoxNumber, self).__init__()
        self.input_shape_custom = self_input_shape
        self.input_shape_accu = 1
        for item in self_input_shape:
            self.input_shape_accu *= item
        
        self.flatten_layer = tf.keras.layers.Flatten()
        self.internal_dense_layer = tf.keras.layers.Dense(1, activation='tanh')

    @tf.function
    def call(self, inputs):
        inputs = self.flatten_layer(inputs)
        output = self.internal_dense_layer(inputs)
        output = tf.divide(output, 2) + 0.5
        return output

2. 调整训练循环的输入格式 #

确保输入到层的张量带有batch维度，即使是单个样本，也要用tf.expand_dims添加batch轴：

for epoch in range(50):
    print("Epoch:", epoch, "of 50")
    average_loss = 0

    for iter, item in enumerate(image):
        num_bounding_boxes = tf.shape(bb[iter])[0]
        float_target = tf.cast(1/num_bounding_boxes, tf.float32)

        with tf.GradientTape() as tape:
            # 给输入添加batch维度（从[H,W,C]变为[1,H,W,C]）
            input_with_batch = tf.expand_dims(item, axis=0)
            logits = dense_bb_num_layer(input_with_batch)
            # 移除batch维度，和目标值匹配
            logits = tf.squeeze(logits, axis=0)
            loss = NumLoss(logits, float_target)
            print("Logits:", logits, "Target:", float_target, "Loss:", loss)
        
        average_loss += loss

        call_gradients = tape.gradient(loss, dense_bb_num_layer.trainable_weights)
        call_optimizer.apply_gradients(zip(call_gradients, dense_bb_num_layer.trainable_weights))

    average_loss /= len(image)
    print(average_loss)

3. 可选：优化损失函数 #

平方损失在输出饱和时梯度极小，可尝试使用平均绝对误差（MAE）替代，提升梯度传递效率：

def NumLoss(logits, expected):
    return tf.abs(logits - expected)

验证方法 #

修改完成后，先单独调用层测试输出是否在0-1之间，再启动训练循环，观察logits是否能输出如0.333、0.25等目标值，损失是否逐步下降。

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