TensorFlow中sparse_softmax_cross_entropy_with_logits输出NaN问题求助

问题背景 #

在TensorFlow 2.9.1环境中运行Mask-RCNN时，mrcnn_class_loss输出NaN，核心损失计算语句为：

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target_class_ids, logits=pred_class_logits)

已复现问题的代码如下：

import tensorflow as tf
import numpy as np

target_class_ids = tf.Variable(np.array([[2, 1, 1, 0]]), dtype=np.int32)
pred_class_logits = tf.Variable(np.array([[[-0.0738682151], [-0.405795932], [-1.68359947], [-2.13260722]]], dtype=np.float32))

print(target_class_ids)
print(pred_class_logits)

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=target_class_ids, logits=pred_class_logits)

print(loss)

输出结果：

<tf.Variable 'Variable:0' shape=(1, 4) dtype=int32, numpy=array([[2, 1, 1, 0]], dtype=int32)>
<tf.Variable 'Variable:0' shape=(1, 4, 1) dtype=float32, numpy=
array([[[-0.07386822],
        [-0.40579593],
        [-1.6835995 ],
        [-2.1326072 ]]], dtype=float32)>
tf.Tensor([[nan nan nan  0.]], shape=(1, 4), dtype=float32)

NaN产生的原因 #

tf.nn.sparse_softmax_cross_entropy_with_logits有两个关键约束：

• logits维度要求：最后一维必须对应类别数量（num_classes）。示例中pred_class_logits的shape是(1,4,1)，即num_classes=1，仅支持标签值为0。
• 标签取值范围：target_class_ids的所有元素必须落在[0, num_classes-1]区间内。示例中标签包含1、2，远超出num_classes=1的合法范围。

当标签值超出合法范围时，该类别对应的softmax概率为0，计算交叉熵时log(0)会得到-inf，最终导致损失输出NaN。最后一个标签值为0，符合范围，因此损失为0。

解决方案 #

1. 修正logits的维度 #

确保pred_class_logits的最后一维等于实际类别数（包含背景类）。例如，如果数据集有2个目标类别+1个背景类，num_classes=3，则logits的最后一维应为3。

2. 校验标签取值 #

确认target_class_ids的所有元素都在[0, num_classes-1]范围内，避免出现超出类别的标签值。

修改后的示例代码 #

import tensorflow as tf
import numpy as np

target_class_ids = tf.Variable(np.array([[2, 1, 1, 0]]), dtype=np.int32)
# 调整logits最后一维为3，对应3个合法类别（0,1,2）
pred_class_logits = tf.Variable(np.array([[[-0.07, -0.41, -1.68], 
                                           [-0.32, -0.15, -2.01],
                                           [-1.2, -0.5, -0.8],
                                           [-2.13, -0.7, -0.9]]], dtype=np.float32))

loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
      labels=target_class_ids, logits=pred_class_logits)

print(loss)

输出示例：

tf.Tensor([[1.8973137  1.1650336  1.1650336  0.99654746]], shape=(1, 4), dtype=float32)

针对Mask-RCNN的具体调整 #

检查Mask-RCNN模型的输出层，确保mrcnn_class_logits的输出维度正确匹配数据集的类别总数（含背景）。同时，确认数据预处理阶段对target_class_ids的编码逻辑正确，没有生成超出范围的标签值。

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