如何在TensorFlow序列模型中实现图像Gradient Difference Loss（仿PyTorch）

实现逻辑 #

完全对齐目标PyTorch仓库的GDL计算逻辑：通过卷积提取图像水平/垂直方向梯度，再计算梯度差异的L1或L2损失，最后按指定方式归约损失值。

代码实现 #

import tensorflow as tf

def gradient_difference_loss(y_true, y_pred, loss_type='l1', reduction='mean'):
    # 获取输入图像的通道数
    channels = y_true.shape[-1]
    
    # 定义水平方向梯度卷积核（适配TensorFlow卷积维度格式）
    kernel_h = tf.constant([[1.0, -1.0]], dtype=tf.float32)
    kernel_h = tf.expand_dims(tf.expand_dims(kernel_h, axis=-1), axis=-1)
    kernel_h = tf.tile(kernel_h, [1, 1, channels, 1])
    
    # 定义垂直方向梯度卷积核
    kernel_v = tf.constant([[1.0], [-1.0]], dtype=tf.float32)
    kernel_v = tf.expand_dims(tf.expand_dims(kernel_v, axis=-1), axis=-1)
    kernel_v = tf.tile(kernel_v, [1, 1, channels, 1])
    
    # 计算真实图像与预测图像的水平/垂直梯度
    grad_true_h = tf.nn.conv2d(y_true, kernel_h, strides=[1,1,1,1], padding='VALID')
    grad_pred_h = tf.nn.conv2d(y_pred, kernel_h, strides=[1,1,1,1], padding='VALID')
    
    grad_true_v = tf.nn.conv2d(y_true, kernel_v, strides=[1,1,1,1], padding='VALID')
    grad_pred_v = tf.nn.conv2d(y_pred, kernel_v, strides=[1,1,1,1], padding='VALID')
    
    # 计算梯度差异
    diff_h = grad_true_h - grad_pred_h
    diff_v = grad_true_v - grad_pred_v
    
    # 选择损失类型
    if loss_type == 'l1':
        loss_h = tf.abs(diff_h)
        loss_v = tf.abs(diff_v)
    elif loss_type == 'l2':
        loss_h = tf.square(diff_h)
        loss_v = tf.square(diff_v)
    else:
        raise ValueError("loss_type must be 'l1' or 'l2'")
    
    # 选择损失归约方式
    if reduction == 'mean':
        loss = tf.reduce_mean(loss_h + loss_v)
    elif reduction == 'sum':
        loss = tf.reduce_sum(loss_h + loss_v)
    elif reduction == 'none':
        loss = loss_h + loss_v
    else:
        raise ValueError("reduction must be 'mean', 'sum' or 'none'")
    
    return loss

序列模型中使用示例 #

构建并编译序列模型时，直接将自定义损失函数传入即可：

# 示例序列模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(256,256,3)),
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Conv2D(64, (3,3), activation='relu'),
    tf.keras.layers.UpSampling2D((2,2)),
    tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
    tf.keras.layers.UpSampling2D((2,2)),
    tf.keras.layers.Conv2D(3, (3,3), activation='sigmoid', padding='same')
])

# 编译模型时指定自定义损失
model.compile(optimizer='adam', loss=gradient_difference_loss)

一致性说明 #

• 卷积核参数、padding方式（VALID对应PyTorch的padding=0）完全匹配目标仓库实现
• 支持L1/L2损失类型及mean/sum/none三种归约方式，与PyTorch版本逻辑一致
• 自动适配多通道图像输入，处理逻辑对齐PyTorch的通道维度操作

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