自定义损失函数致模型val_loss为NaN问题求助

问题背景 #

• 二分类任务，采用自定义损失函数，标签为分类格式，模型最终激活函数为Softmax
• 训练过程中训练损失(loss)正常下降，训练与验证精度(acc)持续提升，但验证损失(val_loss)显示为NaN/inf
• 训练结束后调用model.evaluate(X_train,Y_train)时，训练损失与验证损失均为NaN，但直接调用自定义损失函数计算可得到有效数值
• 已尝试增大batch size、数据缩放、优化器中使用clipnorm/clipvalue等方法，均无效果；模型预测值(y_pred)未出现NaN，怀疑是预测值中的极值导致问题

原自定义损失函数代码 #

def custom_loss(y_true, y_pred):
  import tensorflow as tf
  bce = tf.keras.losses.BinaryCrossentropy(
      from_logits=False,
      label_smoothing=0.0,
      axis=-1,
      reduction="none",
      name="binary_crossentropy",
  )
  intra               = tf.constant(1, dtype=tf.float64)
  inter               = tf.constant(0.01, dtype=tf.float64)

  zeros              = tf.gather_nd(y_pred,tf.where(tf.argmin(y_true, axis = 1)))
  ones               = tf.gather_nd(y_pred,tf.where(tf.argmax(y_true, axis = 1)))

  centroid_zero       = tf.reduce_mean(zeros,0)
  centroid_one        = tf.reduce_mean(ones,0)

  loss_zero_intra     = tf.math.squared_difference(zeros,centroid_zero)
  loss_one_intra      = tf.math.squared_difference(ones,centroid_zero)
  
  loss_zero_intra     = tf.cast(loss_zero_intra, tf.float64)
  loss_one_intra      = tf.cast(loss_one_intra, tf.float64)
  
  loss_intra          = tf.zeros_like(y_pred, tf.float64)
  loss_intra          = tf.tensor_scatter_nd_update(loss_intra,tf.where(tf.argmin(y_true, axis = 1)),loss_zero_intra)
  loss_intra          = tf.tensor_scatter_nd_update(loss_intra,tf.where(tf.argmax(y_true, axis = 1)),loss_one_intra)

  loss_inter_value    = tf.math.sqrt(tf.math.squared_difference(centroid_zero[0],centroid_one[0]) +
                                    tf.math.squared_difference(centroid_zero[1],centroid_one[1]))
  loss_inter = tf.fill(tf.shape(y_pred),loss_inter_value)
  binary_cross_entropy= tf.tile(tf.expand_dims(bce(y_true,y_pred),axis=1),
                                tf.constant([1,2],tf.int32))
  
  loss_intra          = tf.cast(loss_intra, tf.float64)
  loss_inter          = tf.cast(loss_inter, tf.float64)
  binary_cross_entropy= tf.cast(binary_cross_entropy, tf.float64)

  loss                = tf.math.multiply(intra, loss_intra) - tf.math.multiply(inter, loss_inter) + binary_cross_entropy

  return loss

原模型代码 #

def create_model(kernelLength = 32, nb_classes = 2, Chans = 19, Samples = 512,   
                dropoutRate = 0.5 , F1 = 8, D = 2, F2 = 16, norm_rate = 0.25,
                dropoutType = 'Dropout', optimizer_type = 'Adam', lr=0.0005, **kwargs):
  
  K.clear_session()
  gc.collect()
  if dropoutType == 'SpatialDropout2D':
      dropoutType = SpatialDropout2D
  elif dropoutType == 'Dropout':
      dropoutType = Dropout
  else:
      raise ValueError('dropoutType must be one of SpatialDropout2D '
                        'or Dropout, passed as a string.')
  
  input1   = Input(shape = (1, Chans, Samples))

  block1       = Conv2D(F1, (1, kernelLength), padding = 'same',
                                  input_shape = (1, Chans, Samples),
                                  use_bias = False)(input1)
  block1       = BatchNormalization(axis = 1)(block1)
  block1       = DepthwiseConv2D((Chans, 1), use_bias = False, 
                                  depth_multiplier = D,
                                  depthwise_constraint = max_norm(1.))(block1)
  block1       = BatchNormalization(axis = 1)(block1)
  block1       = Activation('elu')(block1)
  block1       = AveragePooling2D((1, 4))(block1)
  block1       = dropoutType(dropoutRate)(block1)
  
  block2       = SeparableConv2D(F2, (1, 16),
                                  use_bias = False, padding = 'same')(block1)
  block2       = BatchNormalization(axis = 1)(block2)
  block2       = Activation('elu')(block2)
  block2       = AveragePooling2D((1, 8))(block2)
  block2       = dropoutType(dropoutRate)(block2)
      
  flatten      = Flatten(name = 'flatten')(block2)
  
  dense        = Dense(nb_classes, name = 'dense', 
                        kernel_constraint = max_norm(norm_rate))(flatten)
  softmax      = Activation('softmax', name = 'softmax')(dense)
  
  model        = Model(inputs=input1, outputs=softmax)
  if optimizer_type == 'Adam':
    optimizer = Adam(learning_rate = lr)
  if optimizer_type == 'Adamax':
    optimizer = Adamax(learning_rate = lr)
  if optimizer_type == 'AdamW':
    optimizer = AdamW(learning_rate = lr)
  model.compile(loss=custom_loss, optimizer=optimizer, metrics = ['accuracy'])
  return model

问题根源分析 #

1. 损失维度不匹配：自定义损失返回二维张量（每个样本对应2个损失值），但Keras期望损失函数返回每个样本的标量损失，验证阶段的损失聚合逻辑因维度异常产生NaN。
2. 空类别样本处理缺失：当批次中某类样本数量为0时，tf.reduce_mean返回NaN，导致后续计算全部失效。
3. 冗余类型转换：多次float32与float64的转换引发数值精度问题，增加NaN出现概率。
4. 损失项符号风险：- inter * loss_inter会让总损失可能变为负数，反向传播时易引发梯度爆炸或NaN。

解决方案 #

针对上述问题，修改自定义损失函数如下：

修改后的自定义损失函数 #

def custom_loss(y_true, y_pred):
    import tensorflow as tf
    # 适配Softmax输出，改用CategoricalCrossentropy
    ce = tf.keras.losses.CategoricalCrossentropy(
        from_logits=False,
        label_smoothing=0.0,
        axis=-1,
        reduction="none",
        name="categorical_crossentropy",
    )
    intra = tf.constant(1.0, dtype=tf.float32)
    inter = tf.constant(0.01, dtype=tf.float32)
    epsilon = 1e-8  # 防止除以0或开根号为0

    # 获取两类样本的掩码
    class_0_mask = tf.equal(tf.argmax(y_true, axis=1), 0)
    class_1_mask = tf.equal(tf.argmax(y_true, axis=1), 1)

    # 提取两类样本的预测值
    class_0_preds = tf.boolean_mask(y_pred, class_0_mask)
    class_1_preds = tf.boolean_mask(y_pred, class_1_mask)

    # 计算类别中心，空类别时用批次均值替代，避免NaN
    centroid_0 = tf.cond(
        tf.size(class_0_preds) > 0,
        lambda: tf.reduce_mean(class_0_preds, axis=0),
        lambda: tf.reduce_mean(y_pred, axis=0)
    )
    centroid_1 = tf.cond(
        tf.size(class_1_preds) > 0,
        lambda: tf.reduce_mean(class_1_preds, axis=0),
        lambda: tf.reduce_mean(y_pred, axis=0)
    )

    # 计算类内损失（聚合为标量）
    loss_intra_0 = tf.reduce_sum(tf.math.squared_difference(class_0_preds, centroid_0), axis=1)
    loss_intra_1 = tf.reduce_sum(tf.math.squared_difference(class_1_preds, centroid_1), axis=1)

    # 构建批次类内损失张量
    loss_intra = tf.zeros(tf.shape(y_true)[0], dtype=tf.float32)
    loss_intra = tf.tensor_scatter_nd_update(loss_intra, tf.where(class_0_mask), loss_intra_0)
    loss_intra = tf.tensor_scatter_nd_update(loss_intra, tf.where(class_1_mask), loss_intra_1)

    # 计算类间损失：用倒数替代减法，确保损失非负
    loss_inter_value = tf.sqrt(
        tf.math.squared_difference(centroid_0[0], centroid_1[0]) +
        tf.math.squared_difference(centroid_0[1], centroid_1[1]) + epsilon
    )
    loss_inter = tf.fill(tf.shape(loss_intra), inter / (loss_inter_value + epsilon))

    # 交叉熵损失（标量）
    ce_loss = ce(y_true, y_pred)

    # 总损失：所有项均为标量，符合Keras要求
    total_loss = intra * loss_intra + loss_inter + ce_loss

    # 可选：添加数值监控，快速定位异常
    # tf.debugging.check_numerics(total_loss, "Total loss contains NaN/inf")

    return total_loss

额外优化建议 #

• 模型输出用Softmax时，搭配CategoricalCrossentropy更匹配分类格式，避免BinaryCrossentropy的适配问题
• 保持模型所有张量精度一致（默认float32），减少跨精度计算的数值误差
• 训练时可启用tf.debugging.check_numerics，实时捕捉损失或梯度中的NaN/inf

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