如何让Keras CategoricalCrossentropy返回指定形状张量而非单值？

问题根源在于Keras的CategoricalCrossentropy类默认会对损失执行归约操作，将所有非类别维度的损失值合并为一个标量，这就是你得到单值结果的原因。要保留每个空间位置的损失并得到[batch_size,5,5,1]的张量，可通过以下两种方式实现：

方案一：使用函数式API直接计算 #

tf.keras.losses.categorical_crossentropy函数默认不会归约损失，仅对指定的类别维度（默认最后一维）计算交叉熵，得到形状为[batch_size,5,5]的张量，最后手动扩展一个维度即可得到目标形状。

示例代码：

import tensorflow as tf

# 模拟输入张量（batch_size=2为例）
a = tf.random.uniform((2, 5, 5, 10))
b = tf.one_hot(tf.random.randint(0, 10, (2, 5, 5)), depth=10)

# 计算每个空间位置的交叉熵，输出形状[2,5,5]
cce_loss = tf.keras.losses.categorical_crossentropy(a, b)
# 扩展最后一维，得到目标形状[2,5,5,1]
cce_loss = tf.expand_dims(cce_loss, axis=-1)

print(cce_loss.shape)  # 输出: (2, 5, 5, 1)

方案二：设置归约模式为NONE #

实例化CategoricalCrossentropy类时，指定reduction=tf.keras.losses.Reduction.NONE，关闭自动归约功能，同样得到[batch_size,5,5]的损失张量，再扩展维度即可。

示例代码：

import tensorflow as tf

# 模拟输入张量
a = tf.random.uniform((2, 5, 5, 10))
b = tf.one_hot(tf.random.randint(0, 10, (2, 5, 5)), depth=10)

# 实例化时关闭归约
bce = tf.keras.CategoricalCrossentropy(reduction=tf.keras.losses.Reduction.NONE)
cce_loss = bce(a, b)
# 扩展最后一维
cce_loss = tf.expand_dims(cce_loss, axis=-1)

print(cce_loss.shape)  # 输出: (2, 5, 5, 1)

补充说明：你之前调整axis参数仅改变了交叉熵计算的类别维度，但并未影响归约行为，因此仍得到标量结果。只有通过设置reduction=NONE或使用无归约的函数式API，才能保留每个空间位置的损失值。

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