基于Unet的I2I翻译任务中自定义广义Dice Loss无梯度问题求助

问题背景 #

基于配对数据集[输入图像，ground_truth]构建I2I翻译任务的Unet神经网络，输入输出均为64x64x3的RGB图像。计划通过自定义损失函数计算真实图像与生成图像的差异：将二者输入已训练模型得到64x64x3的4分类掩码，再用Generalized Dice Loss计算损失，但运行时出现无梯度报错。

报错信息 #

ValueError: No gradients provided for any variable:

完整报错栈：

ValueError: in user code:

    File "/var/folders/6z/93vfmb1930j2ys97s5xf8tp00000gn/T/ipykernel_61471/2960572911.py", line 11, in train_step  *
        generator_optimizer.apply_gradients(zip(generator_gradients,
    File "/usr/local/lib/python3.11/site-packages/keras/src/optimizers/optimizer.py", line 1222, in apply_gradients  **
        grads_and_vars = self.aggregate_gradients(grads_and_vars)
    File "/usr/local/lib/python3.11/site-packages/keras/src/optimizers/optimizer.py", line 1184, in aggregate_gradients
        return optimizer_utils.all_reduce_sum_gradients(grads_and_vars)
    File "/usr/local/lib/python3.11/site-packages/keras/src/optimizers/utils.py", line 33, in all_reduce_sum_gradients
        filtered_grads_and_vars = filter_empty_gradients(grads_and_vars)
    File "/usr/local/lib/python3.11/site-packages/keras/src/optimizers/utils.py", line 77, in filter_empty_gradients
        raise ValueError(

    ValueError: No gradients provided for any variable:

现有代码 #

损失函数代码 #

def dice_coef_4cat(y_true, y_pred, smooth=1e-7):
    y_true_f = K.flatten(K.one_hot(K.cast(y_true, 'int32'), num_classes=4)[...,1:])
    y_pred_f = K.flatten(K.one_hot(K.cast(y_pred, 'int32'), num_classes=4)[...,1:])
    intersect = K.sum(y_true_f * y_pred_f, axis=-1)
    denom = K.sum(y_true_f + y_pred_f, axis=-1)
    
    return K.mean((2. * intersect / (denom + smooth)))

def dice_coef_4cat_loss(y_true, y_pred): 
    return 1 - dice_coef_4cat(y_true, y_pred)

def generator_loss(target, gen_output):
    target_mask = trained_model(target)
    generated_mask = trained_model(gen_output)

    return dice_coef_4cat_loss(target_mask,generated_mask)

训练步骤代码 #

@tf.function
def train_step(input_image, target, step):
    with tf.GradientTape() as gen_tape :
        gen_output = generator(input_image, training=True)

        gen_loss = dice_coef_4cat_loss(target, gen_output)

    generator_gradients = gen_tape.gradient(gen_loss,
                                          generator.trainable_variables)
    
    generator_optimizer.apply_gradients(zip(generator_gradients,
                                          generator.trainable_variables))

问题分析与解决方案 #

核心问题 #

1. 训练步骤未调用正确损失函数：当前train_step直接使用dice_coef_4cat_loss(target, gen_output)，跳过了预训练模型生成掩码的逻辑，未按预期计算掩码间的损失。
2. 离散操作阻断梯度：K.cast(y_pred, 'int32')将连续的生成结果转为整数，属于不可导操作，直接切断梯度回传路径。预训练模型输出的掩码应为概率分布，无需转成整数。

修正后的代码 #

修正损失函数

def dice_coef_4cat(y_true, y_pred, smooth=1e-7):
    # 直接使用预训练模型输出的概率分布，无需转整数再one-hot
    y_true_f = K.flatten(y_true[...,1:])  # 若无需背景类则保留此切片
    y_pred_f = K.flatten(y_pred[...,1:])
    intersect = K.sum(y_true_f * y_pred_f, axis=-1)
    denom = K.sum(y_true_f + y_pred_f, axis=-1)
    
    return K.mean((2. * intersect / (denom + smooth)))

def dice_coef_4cat_loss(y_true, y_pred): 
    return 1 - dice_coef_4cat(y_true, y_pred)

def generator_loss(target, gen_output):
    # 锁定预训练模型参数，避免被更新
    trained_model.trainable = False
    # 推理模式下运行预训练模型
    target_mask = trained_model(target, training=False)
    generated_mask = trained_model(gen_output, training=False)

    return dice_coef_4cat_loss(target_mask, generated_mask)

修正训练步骤

@tf.function
def train_step(input_image, target, step):
    with tf.GradientTape() as gen_tape :
        gen_output = generator(input_image, training=True)
        # 调用正确的掩码损失函数
        gen_loss = generator_loss(target, gen_output)

    generator_gradients = gen_tape.gradient(gen_loss, generator.trainable_variables)
    generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))

额外注意事项 #

• 若预训练模型输出为logits，需先通过tf.nn.softmax转为概率分布后再传入损失函数。
• 确认target_mask的生成逻辑符合任务需求：当前代码用trained_model(target)生成真实掩码，需保证该逻辑与任务目标一致。

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