TensorFlow自定义损失函数：实现L1与MS-SSIM混合损失

核心实现步骤 #

先实现可微分的高斯滤波逻辑，再按公式结合MS-SSIM损失与滤波后的L1损失。

1. 高斯滤波函数实现

针对图像张量的每个通道独立应用高斯滤波，保证可微分以适配TensorFlow自动求导：

import tensorflow as tf

def gaussian_filter(input_tensor, kernel_size=11, sigma=1.5):
    # 生成1D高斯核
    x = tf.range(-kernel_size//2 + 1, kernel_size//2 + 1, dtype=tf.float32)
    gauss_1d = tf.exp(-(x**2) / (2 * sigma**2))
    gauss_1d /= tf.reduce_sum(gauss_1d)
    
    # 扩展为2D高斯核
    gauss_2d = tf.tensordot(gauss_1d, gauss_1d, axes=0)
    
    # 适配depthwise卷积的形状：(kernel_size, kernel_size, channels, 1)
    channels = tf.shape(input_tensor)[-1]
    gauss_kernel = tf.tile(
        tf.expand_dims(tf.expand_dims(gauss_2d, -1), -1),
        [1, 1, channels, 1]
    )
    
    # 应用滤波，保持输入输出尺寸一致
    filtered_tensor = tf.nn.depthwise_conv2d(
        input_tensor, gauss_kernel, strides=[1,1,1,1], padding='SAME'
    )
    return filtered_tensor

2. MS-SSIM损失函数

基于TensorFlow内置的多尺度SSIM接口，转换为损失形式：

def ms_ssim_loss(y_true, y_pred):
    # 计算每个样本的MS-SSIM值，取batch平均后转为损失
    ms_ssim = tf.image.ssim_multiscale(y_true, y_pred, max_val=1.0)
    return 1 - tf.reduce_mean(ms_ssim)

3. 完整混合损失函数

按公式组合两种损失：

def mix_loss(y_true, y_pred):
    alpha = 0.84  # 权重参数可根据需求调整
    
    # 计算MS-SSIM损失
    ms_ssim_l = ms_ssim_loss(y_true, y_pred)
    
    # 计算L1误差图，滤波后取平均得到损失
    l1_error_map = tf.abs(y_true - y_pred)
    filtered_l1_error = gaussian_filter(l1_error_map)
    filtered_l1_l = tf.reduce_mean(filtered_l1_error)
    
    # 组合最终损失
    return alpha * ms_ssim_l + (1 - alpha) * filtered_l1_l

关键注意事项 #

• 输入张量y_true和y_pred需为浮点数类型，若图像取值范围是[0,255]，需先归一化到[0,1]，或调整tf.image.ssim_multiscale的max_val参数为255
• 高斯核的kernel_size和sigma可参考原论文或你提供的Caffe代码调整，示例参数与Caffe代码一致
• 在Keras训练流程中，可直接将mix_loss作为自定义损失函数传入model.compile()

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