基于ResNet50的图像分割模型正则化损失出现NaN的解决咨询

问题背景 #

基于ResNet50编码器+带反池化/跳跃连接的解码器实现图像分割任务，损失计算代码如下：

total_loss = tf.add_n([dice_coefficient_output+IOU_output]+tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES))

训练第1轮总损失约0.4，第2轮出现NaN，排查确认tf.GraphKeys.REGULARIZATION_LOSSES中多数层返回NaN。已尝试图像归一化（0-1/-1-1/Z-score）、降低学习率、调整L2权重衰减、减少神经元数量等手段，仅能推迟NaN出现时间至第4轮。

优化建议 #

• 精准控制正则化应用范围：L2正则化仅对卷积/全连接层的kernel参数生效，禁止应用在偏置项、BN层的gamma/beta参数上——这类参数的数值波动极易引发NaN。手动指定正则化对象：

  kernel_reg = tf.contrib.layers.l2_regularizer(scale=1e-4)
  conv = tf.layers.conv2d(inputs, filters, kernel_size, kernel_regularizer=kernel_reg)
  

• 手动实现正则化损失并截断参数：放弃自动收集的REGULARIZATION_LOSSES，手动遍历目标参数计算L2损失，同时限制参数数值范围避免溢出：

  reg_loss = 0.0
  l2_scale = 1e-4
  for var in tf.trainable_variables():
      if 'kernel' in var.name:
          clipped_var = tf.clip_by_value(var, -10.0, 10.0)
          reg_loss += tf.nn.l2_loss(clipped_var) * l2_scale
  total_loss = dice_coefficient_output + IOU_output + reg_loss
  

• 修复Dice/IOU的数值稳定性：计算时添加极小平滑项（如1e-7）防止除以0，避免反向传播梯度爆炸影响参数更新：

  def dice_coeff(y_true, y_pred):
      smooth = 1e-7
      inter = tf.reduce_sum(y_true * y_pred)
      return (2. * inter + smooth) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + smooth)
  

• 启用全局梯度裁剪：限制梯度最大范数，避免参数更新幅度过大导致数值异常：

  opt = tf.train.AdamOptimizer(learning_rate=1e-5)
  grads_vars = opt.compute_gradients(total_loss)
  clipped_grads = [(tf.clip_by_norm(g, 5.0), v) for g, v in grads_vars if g is not None]
  train_op = opt.apply_gradients(clipped_grads)
  

• 检查反池化层实现：若使用自定义反池化逻辑，需确保索引计算无错误，可添加数值截断操作过滤异常值：

  unpooled = tf.nn.max_unpool_with_argmax(pooled, argmax, ksize=[1,2,2,1], strides=[1,2,2,1], output_shape=orig_shape)
  unpooled = tf.clip_by_value(unpooled, -1e3, 1e3)
  

• 监控数值变化定位问题：添加日志或TensorBoard监控，记录每轮正则化损失、各层参数的均值/最大值，定位首个出现NaN的参数层：

  tf.summary.scalar('reg_loss', reg_loss)
  for var in tf.trainable_variables():
      tf.summary.histogram(var.name, var)
  

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