如何在TensorFlow中计算Softmax输出的期望值？

针对你的需求——基于有序数值类别的Softmax输出计算期望值（$\sum_{i=0}^{N-1} i \cdot p_i$），这里提供两种实现方案：

一、将期望值作为模型的额外输出层实现 #

可以自定义一个固定权重的层，直接基于Softmax输出计算期望值，且该层不参与训练。

自定义期望值层 #

import tensorflow as tf
from tensorflow.keras.layers import Layer

class ExpectedValueLayer(Layer):
    def __init__(self, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.num_classes = num_classes
        # 初始化固定权重：形状为(N,1)，值为0到N-1的浮点数
        self.fixed_weights = tf.Variable(
            initial_value=tf.range(num_classes, dtype=tf.float32)[:, tf.newaxis],
            trainable=False  # 固定权重，不参与反向传播
        )
    
    def call(self, softmax_output):
        # 矩阵乘法计算期望值：(batch_size, N) @ (N,1) → (batch_size,1)
        return tf.matmul(softmax_output, self.fixed_weights)

整合到模型中 #

假设你已有基础模型，最后一层是Softmax输出，只需添加这个自定义层即可得到多输出模型：

# 示例基础模型
N = 5  # 类别数量
base_model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(N, activation='softmax', name='softmax_probs')
])

# 添加期望值输出层
expected_value = ExpectedValueLayer(N)(base_model.output)
# 构建多输出模型
multi_output_model = tf.keras.Model(
    inputs=base_model.input,
    outputs=[base_model.output, expected_value]
)

此时模型会输出两个结果：Softmax概率分布，以及对应的1维期望值，后续可以对期望值单独定义损失函数。

二、在自定义损失类中实现期望值计算 #

如果不需要将期望值作为模型输出，而是直接在损失计算中使用，可以在自定义损失类里完成期望值的计算逻辑。

自定义损失类 #

class OrderedClassLoss(tf.keras.losses.Loss):
    def __init__(self, num_classes, **kwargs):
        super().__init__(**kwargs)
        self.num_classes = num_classes
        # 生成类别索引张量：形状(1, N)，值为0到N-1，用于广播计算
        self.class_indices = tf.range(num_classes, dtype=tf.float32)[tf.newaxis, :]
    
    def call(self, y_true, y_pred):
        # y_pred是Softmax输出，形状(batch_size, N)
        # 计算每个样本的期望值：逐元素相乘后沿类别轴求和
        expected_value = tf.reduce_sum(y_pred * self.class_indices, axis=1)
        
        # 这里替换为你需要的损失逻辑，比如以MSE为例（假设y_true是数值型真实标签）
        loss = tf.keras.losses.mean_squared_error(y_true, expected_value)
        return tf.reduce_mean(loss)

使用自定义损失 #

直接将该损失类传入模型的compile方法即可：

model.compile(optimizer='adam', loss=OrderedClassLoss(N))

关键说明 #

生成类别索引时，tf.range(num_classes)会直接创建0到N-1的序列，通过tf.newaxis调整形状后，能和Softmax输出（形状(batch_size, N)）实现广播相乘，无需额外的ones_like类方法。

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