Keras中U-Net自定义损失函数y_true形状异常问题求助

我来帮你分析下这个问题，核心原因是你把batch_size作为参数传入了损失函数，这打乱了Keras对张量形状的自动推断逻辑，导致y_true的形状被错误地标记为(None, None, None, None)。下面是具体的解决思路和修正后的代码：

问题根源 #

Keras在构建损失函数时，会依赖动态张量形状来自动适配不同的批量大小，而你硬编码batch_size后，Keras无法正确推断y_true的实际维度（比如(batch_size, 128, 256, 3)），只能用None填充所有维度。另外，你的梯度计算函数里手动指定batch_size切片也是多余的，Keras的张量操作本身就支持批量维度的自动处理。

修正方案 #

1. 移除损失函数中的batch_size参数 #

不需要手动传入batch_size，Keras会自动处理任意批量大小的张量。

2. 简化梯度计算逻辑 #

直接利用张量切片操作（[:, 1:, ...]和[:, :-1, ...]）来计算梯度，不需要手动初始化零张量或者指定批量范围。

修正后的完整代码 #

import tensorflow.keras.backend as K

# Encouraging the predicted image to match the label not only in image domain, but also in gradient domain
def keras_customized_loss(lambda1=1.0, lambda2=0.05):
    def grad_x(image):
        # 计算x方向的梯度（水平方向相邻像素差的绝对值）
        return K.abs(image[:, 1:, :, :] - image[:, :-1, :, :])
    
    def grad_y(image):
        # 计算y方向的梯度（垂直方向相邻像素差的绝对值）
        return K.abs(image[:, :, 1:, :] - image[:, :, :-1, :])
    
    def compute_loss(y_true, y_pred):
        # 计算预测和真实图像的梯度
        pred_grad_x = grad_x(y_pred)
        pred_grad_y = grad_y(y_pred)
        true_grad_x = grad_x(y_true)
        true_grad_y = grad_y(y_true)
        
        # 计算各项损失
        mse_loss = K.mean(K.square(y_pred - y_true))
        grad_x_loss = K.mean(K.square(pred_grad_x - true_grad_x))
        grad_y_loss = K.mean(K.square(pred_grad_y - true_grad_y))
        
        # 加权合并损失
        return lambda1 * mse_loss + lambda2 * grad_x_loss + lambda2 * grad_y_loss
    
    return compute_loss

# 编译模型时不需要传入batch_size
model.compile(optimizer='adam', loss=keras_customized_loss(), metrics=['MeanAbsoluteError'])

额外注意事项 #

• 确保你的训练数据生成器（或者输入数据）输出的y_true形状确实是(batch_size, 128, 256, 3)，和模型的输出形状完全匹配。
• 如果后续需要验证张量形状，可以在compute_loss函数中临时添加print(K.int_shape(y_true))来查看实际形状，确认问题是否解决。

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