TensorFlow训练时如何在TensorBoard中显示各分类的F1分数

问题原因 #

TensorFlow 2.5的内置TensorBoard回调默认会将返回数组形式的指标（即你设置average=None得到的多类别F1数组）自动聚合为单个标量，因此训练过程中只会显示一条汇总曲线，不会拆分单类指标。你可以通过以下两种方案实现单类F1分数的TensorBoard可视化：

方案1：自定义回调（推荐，适配任意类别数量） #

你可以自定义训练回调，在每个epoch结束时单独计算每个类别的F1值并写入TensorBoard日志，代码示例如下：

import tensorflow as tf
from tensorflow.keras.callbacks import Callback

class PerClassF1Callback(Callback):
    def __init__(self, eval_dataset, class_names=None, log_dir='./logs'):
        super().__init__()
        self.eval_data = eval_dataset # 可以传训练集或验证集
        # 未传入类别名时自动生成默认名称
        self.class_names = class_names if class_names else [f'class_{i}' for i in range(CLASSES)]
        self.writer = tf.summary.create_file_writer(log_dir + '/per_class_f1')
        # 初始化F1指标，和你编译模型时的参数保持一致
        self.f1_metric = tfa.metrics.F1Score(num_classes=CLASSES, average=None)
    
    def on_epoch_end(self, epoch, logs=None):
        # 重置指标状态避免历史数据干扰
        self.f1_metric.reset_states()
        # 遍历数据集计算当前epoch的F1分数
        for x, y_true in self.eval_data:
            y_pred = self.model.predict(x, verbose=0)
            self.f1_metric.update_state(y_true, y_pred)
        f1_scores = self.f1_metric.result().numpy()
        # 逐个类别写入TensorBoard日志
        with self.writer.as_default():
            for idx, score in enumerate(f1_scores):
                tf.summary.scalar(f'f1_score/{self.class_names[idx]}', score, step=epoch)
        self.writer.flush()
        # 可选：将单类F1加入日志，训练时终端也会打印
        for idx, score in enumerate(f1_scores):
            logs[f'f1_{self.class_names[idx]}'] = score

回调使用方式 #

在model.fit时将自定义回调传入回调列表即可：

# 初始化回调，这里以验证集为例，也可以传入训练集
f1_callback = PerClassF1Callback(eval_dataset=val_dataset, log_dir='./logs')
# 保留原有的TensorBoard回调
tb_callback = tf.keras.callbacks.TensorBoard(log_dir='./logs')

model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=你的训练轮数,
    callbacks=[tb_callback, f1_callback]
)

方案2：拆分单类指标（适合类别数量少的场景） #

你可以在编译模型时，为每个类别单独定义一个F1指标，内置的TensorBoard回调会自动识别为独立标量生成单独曲线：

metrics = [
    # 保留原有的整体多类F1指标
    tfa.metrics.F1Score(num_classes=CLASSES, average=None, name='overall_f1'),
    # 为每个类别单独定义F1指标，指定对应的class_id
    tfa.metrics.F1Score(num_classes=CLASSES, average='macro', class_id=0, name='f1_class_0'),
    tfa.metrics.F1Score(num_classes=CLASSES, average='macro', class_id=1, name='f1_class_1'),
    # 有多少个类就加多少个，依次类推
    # tfa.metrics.F1Score(num_classes=CLASSES, average='macro', class_id=N, name='f1_class_N')
]
model.compile(
    optimizer='adam',
    loss=tf.keras.losses.CategoricalCrossentropy(),
    metrics=metrics
)

该方案无需额外自定义逻辑，缺点是类别数量较多时配置重复度较高。

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