PyTorch Lightning如何在控制台输出训练运行汇总表格

配置完成后不影响原有TensorBoard日志的正常写入，控制台可输出与HuggingFace Transformers格式一致的逐epoch指标汇总表。

具体配置步骤 #

• 自定义进度条回调，过滤step级冗余日志，只保留epoch维度的聚合指标
• 初始化Trainer时关闭逐step日志打印，绑定自定义进度条与TensorBoard Logger
• 在模型的训练、验证步骤中，配置self.log按epoch维度记录需要展示的指标

可直接复用的代码示例 #

from pytorch_lightning import Trainer, LightningModule
from pytorch_lightning.callbacks import TQDMProgressBar
from pytorch_lightning.loggers import TensorBoardLogger

# 自定义进度条，适配逐epoch表格输出
class EpochTableProgressBar(TQDMProgressBar):
    def get_metrics(self, trainer, model):
        metric_map = super().get_metrics(trainer, model)
        # 移除不需要在表格中展示的冗余字段
        for turndown_key in ["v_num", "loss_step", "step"]:
            metric_map.pop(turndown_key, None)
        return metric_map

    def on_train_epoch_end(self, trainer, pl_module):
        super().on_train_epoch_end(trainer, pl_module)
        # 每个epoch输出后换行，保证表格行对齐
        print()

# 初始化TensorBoard日志器，正常写入日志供TensorBoard读取
tb_logger = TensorBoardLogger(save_dir="./training_logs")

# 初始化训练器
trainer = Trainer(
    logger=tb_logger,
    callbacks=[EpochTableProgressBar(leave=True)],
    max_epochs=10, # 按实际训练需求修改
    log_every_n_steps=10**9, # 屏蔽逐step的日志输出
    enable_model_summary=False, # 关闭初始模型结构打印，避免打乱表格格式
    benchmark=True, # 自动统计运行耗时、样本吞吐率指标
)

在自定义的LightningModule实现中，记录指标时需要指定epoch级聚合，同时开启进度条展示权限：

class YourTrainModel(LightningModule):
    def training_step(self, batch, batch_idx):
        # 编写训练逻辑，计算训练损失
        train_loss = ...
        # 记录训练损失，仅在epoch维度聚合，同步到进度条表格
        self.log("Training Loss", train_loss, on_step=False, on_epoch=True, prog_bar=True)
        return train_loss
    
    def validation_step(self, batch, batch_idx):
        # 编写验证逻辑，计算验证损失
        val_loss = ...
        # 记录验证损失
        self.log("Validation Loss", val_loss, on_step=False, on_epoch=True, prog_bar=True)

补充说明 #

• 上述配置基于PyTorch Lightning 1.8及以上版本验证可用，2.x版本无需额外调整参数
• 如果想要更美观的带格式表格，可以将TQDMProgressBar替换为内置的RichProgressBar，配置逻辑完全一致，输出的表格会自动做列对齐、样式优化
• 运行时间、每秒样本数两个指标为框架自动统计，不需要手动编写代码记录

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