TensorFlow自定义训练循环中按类别计算Macro平均的问题

你的问题核心是图模式下不能用Python原生循环/列表操作遍历或动态修改张量，因为图构建阶段需要静态定义所有操作，而Python循环是在图构建时执行，无法适配张量的动态维度。下面是适配图模式的改写方案：

核心思路 #

用TensorFlow原生的tf.map_fn替代Python循环，实现对每个分类任务（loss的列）的批量处理；同时在单任务处理中，用tf.boolean_mask替代tf.where+tf.gather_nd的组合，简化操作并适配图模式。

改写后的代码 #

def compute_macro_avg_per_class(loss_col, y_col):
    # 获取当前分类任务下的唯一标签值
    unique_labels, _ = tf.unique(y_col)
    
    # 对每个唯一标签，计算对应loss的平均值
    per_label_avg = tf.map_fn(
        lambda label: tf.reduce_mean(tf.boolean_mask(loss_col, tf.equal(y_col, label))),
        unique_labels,
        dtype=tf.float32
    )
    
    # 对所有标签的loss平均值再求平均（macro平均）
    return tf.reduce_mean(per_label_avg)

# 对每个分类任务（loss的每一列）应用上述函数
items_loss_list = tf.map_fn(
    lambda idx: compute_macro_avg_per_class(loss[:, idx], y[:, idx]),
    tf.range(tf.shape(loss)[1]),
    dtype=tf.float32
)

关键改动说明 #

1. 替换Python循环为tf.map_fn：tf.map_fn是TensorFlow图模式支持的迭代操作，会将逻辑转化为图节点，而非Python层面的循环。
2. 用tf.boolean_mask简化索引：相比tf.where+tf.gather_nd，tf.boolean_mask更简洁，且天然支持图模式下的掩码操作。
3. 避免修改Python列表：图模式下不能通过append修改Python列表，所有结果收集都要通过TensorFlow张量操作完成。

为什么你的tf.while_loop尝试失败？ #

你在while_loop的body函数中修改Python列表item_score_avg，这属于图构建阶段的Python侧操作，无法被TensorFlow图捕获。图模式下必须用TensorFlow的张量变量（如tf.TensorArray）来动态收集循环结果，而tf.map_fn已经封装了这种逻辑，无需手动实现while循环。

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