PyTorch Lightning如何将训练/验证损失接入TensorBoard HParams

在你当前使用的PyTorch Lightning 1.2.6、TensorBoard 2.4.1版本组合下，出现该现象是两个明确的适配逻辑导致的：

• 框架默认只会将预留的hp_metric指标纳入HParams面板的可对比指标列表，不会自动把你在训练步、验证步通过self.log手动记录的损失和HParams超参配置建立绑定关系
• TensorBoard扫描日志事件文件时，只有被显式登记为HParams关联项的指标，才会出现在左侧边栏的Metrics分类下。未登记的指标即使在SCALARS页能正常绘制曲线，也不会进入HParams的可选指标池。

1. 显式初始化TensorBoard日志器，关闭冗余的默认指标 #

不要使用PyTorch Lightning隐式创建的默认日志器，在训练器初始化阶段手动实例化TB日志器，并关闭默认生成的hp_metric指标：

from pytorch_lightning.loggers import TensorBoardLogger

# 训练器初始化前创建日志器实例
tb_logger = TensorBoardLogger(
    save_dir="tb_logs",
    name="your_experiment_name",
    # 关闭默认的冗余hp_metric，避免列表出现无用条目
    default_hp_metric=False
)

# 原有DDP、GPU配置保持不变
trainer = pl.Trainer(
    logger=tb_logger,
    accelerator="ddp",
    gpus=2,
    # 其他原有训练配置全部保留
)

2. 在模型中显式绑定HParams和目标指标 #

利用PyTorch Lightning的setup钩子（该钩子会在训练流程启动时、分布式各节点上依次执行），在全局主进程节点上，把超参配置和你需要展示的train_loss、val_loss登记给TB日志器，注意DDP场景下必须加主进程判断，避免多进程重复写日志导致事件文件损坏：

class MyLightningModel(pl.LightningModule):
    def __init__(self, your_hparams_dict):
        super().__init__()
        # 保存当前实验的超参，供后续HParams日志写入
        self.save_hyperparameters(your_hparams_dict)

    def setup(self, stage: str):
        # 仅在全局主进程执行日志写入，避免DDP多进程写冲突
        if self.trainer.is_global_zero:
            self.logger.log_hyperparams(
                params=self.hparams,
                metrics={
                    # *这里的key必须和后续self.log的指标名完全一致*
                    "val_loss": None,
                    "train_loss": None
                }
            )

    # 原有训练、验证逻辑保留，仅补充DDP分布式聚合的配置
    def training_step(self, batch):
        x, labels = batch
        out = self(x)
        loss = F.mse_loss(out, labels)
        # 补充sync_dist=True 保证DDP下损失是全局聚合值
        self.log("train_loss", loss, sync_dist=True)
        return loss

    def validation_step(self, batch):
        x, labels = batch
        out = self(x)
        loss = F.mse_loss(out, labels)
        self.log("val_loss", loss, sync_dist=True)
        return loss

3. 运行后验证效果 #

• 删掉旧版本实验生成的日志文件夹，避免TB读取旧事件文件造成指标绑定关系失效
• 跑通至少一个完整的训练+验证周期后，刷新TensorBoard页面进入HPARAMS标签页，左侧Metrics列表就会出现train_loss、val_loss选项
• 勾选需要展示的指标后，即可在平行坐标视图中调整坐标轴顺序，用val_loss替代原有默认指标完成不同超参组合的效果对比

• 传入log_hyperparams的指标key，必须和后续self.log的指标名完全一致，大小写、下划线差异都会导致绑定失败，指标无法显示在HParams列表中
• DDP场景下禁止在非主进程执行log_hyperparams写入操作，否则会生成重复的超参条目，造成平行坐标视图的数据点错位
• 如果多卡训练未开启损失的分布式同步（即没有加sync_dist=True），记录的val_loss仅为单卡局部值，平行坐标视图的对比结果会存在偏差。

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