基于Unet的双输出CNN训练验证分用损失及EarlyStopping实现

• 现有双输出Unet结构：单输入对应两个分支输出，分别为形状参数数组b、分割图像C13
• 预期实现逻辑：
  • 训练阶段仅对b的预测结果计算Shape_loss（欧氏距离），作为反向传播的唯一损失依据
  • 验证阶段仅对C13的预测结果计算Dice_loss，作为模型评估指标
  • EarlyStopping完全绑定验证集Dice_loss值触发
• 原方案失效原因：K.in_train_phase在高阶Keras版本中会与多输出模型的计算图逻辑冲突，无法正确切换训练/验证阶段的损失分支；原代码将损失揉入同一个函数，且EarlyStopping错误绑定训练集loss指标，完全未用到验证集结果。

第一步：拆分独立损失函数 #

删除原有Euclidean_loss里的阶段判断逻辑，分别定义两个独立的损失：

# C13分支使用的Dice损失
def dice_loss(y_true, y_pred):
    y_true = tf.cast(y_true, tf.float32)
    numerator = 2 * tf.reduce_sum(y_true * y_pred)
    denominator = tf.reduce_sum(y_true + y_pred)
    return 1 - numerator / denominator

# b分支使用的Shape损失
def shape_loss(y_true, y_pred):
    return K.sqrt(K.sum(K.square(y_pred - y_true), axis=-1))

第二步：自定义模型类分离训练/验证逻辑 #

通过重写Keras的train_step和test_step方法，完全绕开K.in_train_phase的兼容性问题，独立定义两个阶段的损失计算规则，原网络结构完全不需要修改：

class SplitLossUnet(tf.keras.Model):
    def train_step(self, data):
        # 拆分输入和真值，顺序和数据生成器返回的格式保持一致
        x, (y_true_b, y_true_c13) = data
        with tf.GradientTape() as tape:
            y_pred_b, y_pred_c13 = self(x, training=True)
            # 训练阶段仅计算b分支的Shape损失作为反向传播依据
            train_loss = shape_loss(y_true_b, y_pred_b)
        # 梯度更新逻辑保持默认
        grads = tape.gradient(train_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        # 记录训练阶段的Shape损失值
        self.compiled_metrics.update_state([y_true_b, y_true_c13], [y_pred_b, y_pred_c13])
        results = {m.name: m.result() for m in self.metrics}
        results["train_shape_loss"] = train_loss
        return results

    def test_step(self, data):
        # 拆分输入和真值
        x, (y_true_b, y_true_c13) = data
        y_pred_b, y_pred_c13 = self(x, training=False)
        # 验证阶段仅计算C13分支的Dice损失作为评估依据
        val_dice_loss = dice_loss(y_true_c13, y_pred_c13)
        self.compiled_metrics.update_state([y_true_b, y_true_c13], [y_pred_b, y_pred_c13])
        results = {m.name: m.result() for m in self.metrics}
        # 将验证集Dice损失写入val_loss字段，供EarlyStopping监控
        results["val_loss"] = val_dice_loss
        return results

第三步：替换模型初始化流程 #

把原代码里创建普通Model的部分，替换为上面定义的自定义模型类：

# 原网络主干结构保持不变，获取输出张量
input_im = Input((IMAGES_RESOLUTION[0], IMAGES_RESOLUTION[1], 1))
# 这里保留原代码里所有卷积、上采样、跳跃连接逻辑，输出两个张量：r2（b）、c13
r2 = ... # 原代码里b分支的Dot层输出
c13 = ... # 原代码里C13分支的最后一层卷积输出

# 用自定义类创建模型，替代原来的Model()调用
model = SplitLossUnet(inputs=[input_im], outputs=[r2, c13])

# 模型编译，不需要额外指定损失函数，损失逻辑已经在train_step/test_step里定义
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.00001))

第四步：修正EarlyStopping配置 #

把监控指标从训练集loss改为验证集val_loss（也就是我们在test_step里写入的Dice损失值）：

earlystopping = EarlyStopping(
    monitor='val_loss',
    verbose=0,
    min_delta=0.001,
    patience=100,
    mode='min',
    restore_best_weights=True
)

# 训练流程和原有逻辑完全一致
history = model.fit(
    train_gen,
    epochs=Eps,
    verbose=1,
    callbacks=[earlystopping],
    validation_data=valid_gen,
    steps_per_epoch=30,
    validation_steps=5
)

• 直接删除原代码里的model.outputs[0]._uses_learning_phase = True：这是旧版本Keras的私有API，新版本下已失效，无实际作用
• 不要在损失函数里添加任何阶段判断分支：这种写法会导致计算图构建错误，且训练时额外计算不必要的损失值，拖慢速度
• 数据生成器返回的真值顺序必须和两个输出的顺序严格对应：第一个真值数组对应b的标注，第二个真值图像对应C13的标注，否则会出现张量维度不匹配的报错
• 训练过程中梯度仅来自b分支的Shape损失，不会被C13的Dice损失干扰；验证过程中不会额外计算b分支的Shape损失，减少验证阶段的算力消耗

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