Keras自定义损失函数batch_size>1触发尺寸报错的解决方法

• 自定义损失函数完全没有适配动态批量维度，存在三个硬伤：
  1. 对批量张量的reshape操作压平了batch轴，没有保留批量维度。你计算得到的u_field形状是[batch, H*W]，却被硬reshape成[H, W]，当batch_size=1时元素数刚好匹配，batch_size翻倍后元素数是单样本的两倍，直接触发形状不匹配的报错（对应报错里55778 = 2*27889的数值关系）
  2. 硬拿get_shape()拿到的batch_size值当索引访问张量，逻辑完全错误：张量第一维的合法索引范围是0 ~ batch_size-1，拿batch_size当索引属于越界访问，根本无法命中批量内的有效样本
  3. 损失计算只覆盖了单个样本的逻辑，没有对整个batch的结果做聚合，不符合Keras损失函数的输出要求（需要返回一个标量的batch平均损失）

• 第一步：还原场形状时保留batch轴，用动态维度适配批量大小，不要硬写固定尺寸
• 第二步：差分计算直接对全批量张量做切片运算，自动广播到每个样本，取消手动索引单样本的写法
• 第三步：分母加上极小值避免除零，最后对整个batch的损失求均值聚合

def my_forward_div_loss_fn(input_img):
  # 取固定的空间维度值，batch维度为动态值，不做硬编码
  _, img_h, img_w, _ = input_img.shape
  output_dim = img_h * img_w

  def loss(y_true, y_pred):
    # 展平的预测值还原为空间场，-1 自动适配当前batch大小
    u_field = K.abs(y_true - y_pred)
    u_field = K.reshape(u_field, shape=[-1, img_h, img_w])

    # 拆分输入的各个物理场，保留batch维度
    v_field = input_img[..., 0]
    x_field = input_img[..., 1]
    y_field = input_img[..., 2]

    # 逐样本批量计算向前差分，自动遍历整个batch
    forward_dudx = (u_field[:, 1:-1, 2:] - u_field[:, 1:-1, 1:-1]) / (x_field[:, 1:-1, 2:] - x_field[:, 1:-1, 1:-1] + K.epsilon())
    forward_dvdy = (v_field[:, 1:-1, 1:-1] - v_field[:, 2:, 1:-1]) / (y_field[:, :-2, 1:-1] - y_field[:, 2:, 1:-1] + K.epsilon())

    forward_divergence = forward_dudx + forward_dvdy
    # 替换计算产生的NaN，保留原有逻辑
    forward_divergence = tf.where(tf.math.is_nan(forward_divergence), tf.zeros_like(forward_divergence), forward_divergence)
    # 对整个batch的损失求均值，返回标量
    return K.mean(K.square(forward_divergence))
  return loss

• 不要在损失函数内将batch维度的返回值当作固定常数使用，训练过程中最后一个batch可能出现尺寸不匹配的情况，动态维度适配才能兼容所有训练场景
• 所有自定义张量运算必须保留第一维的batch轴，禁止压平批量维度后做单样本逻辑的硬编码
• 除法计算中加入K.epsilon()避免0分母导致的梯度爆炸

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