如何计算Tensorflow中softmax值的期望值？

Question

我有一个模型，其中包含 N 个类别的最终 softmax 层。 这些类别是有序的和数字的，因此计算 softmax 给出的概率分布的统计数据是有意义的。

假设类别值只是一个递增的索引序列。 所以，第一个值为 0，第二个值为 1，等等。我想计算期望值 (\sum_i=0^{N-1} i p_i)。

如何在 Tensorflow 中将其作为附加输出层（维度 1）来实现？也许使用冷冻权重（0 到 N-1）？然后我可以对这个输出值应用损失。

或者，如果不可能将输出作为模型架构的一部分，那么我如何在损失类中实现它？这需要像 tf.ones_like() 这样的东西沿着轴用递增的整数（0 到 N-1）填充张量。

Answer 1

根据我在问题中的想法，我使用保存类别标签值的冻结权重来实现平均值。 在此示例中，这些是 [0, N-1]

probs = Softmax()(logits)

mean = Dense(1, use_bias=False, trainable=False,
             kernel_initializer=tf.keras.initializers.constant(
                 np.arange(N)) )(probs)