在PyTorch Lightning中使用TensorBoard Logger实现多模型指标同图可视化

嘿，我明白你的困扰！刚接触PyTorch Lightning的Logger确实容易懵，尤其是想对比多个模型的指标时。你现在的问题是两个Trainer默认会生成独立的日志版本文件夹，导致TensorBoard把它们当成完全分开的训练任务，没法在同一张图里对比。下面给你两种简单可行的解决方法：

方法一：用TensorBoardLogger指定不同模型名称（推荐） #

这种方法不需要修改模型内部的日志逻辑，只需要给每个Trainer绑定一个带专属名称的Logger实例，让它们的日志存在同一个根目录下的不同子文件夹里，TensorBoard会自动识别并支持同图对比。

修改后的代码示例： #

from pytorch_lightning.loggers import TensorBoardLogger

# 创建两个Logger，指定相同的保存目录，不同的模型名称
logger1 = TensorBoardLogger("lightning_logs", name="Model1")
logger2 = TensorBoardLogger("lightning_logs", name="Model2")

# 在初始化Trainer时传入对应的logger
trainer1 = pl.Trainer(
    gpus=n_gpus, 
    max_epochs=n_epochs, 
    progress_bar_refresh_rate=20, 
    num_sanity_val_steps=0,
    logger=logger1  # 绑定Model1的Logger
)
trainer2 = pl.Trainer(
    gpus=n_gpus, 
    max_epochs=n_epochs, 
    progress_bar_refresh_rate=20, 
    num_sanity_val_steps=0,
    logger=logger2  # 绑定Model2的Logger
)

# 正常训练
trainer1.fit(Model1, train_loader, val_loader)
trainer2.fit(Model2, train_loader, val_loader)

# 启动TensorBoard
%load_ext tensorboard
%tensorboard --logdir lightning_logs/

查看对比图表的操作： #

启动TensorBoard后，进入Scalars页面：

• 在左侧的标签栏里，你会看到Model1和Model2两个分类，展开后能看到各自的train_loss、val_acc等指标
• 同时勾选两个模型的同一类指标（比如都选val_acc），TensorBoard就会把它们绘制在同一张图表里，方便你直接对比

方法二：手动给指标添加模型前缀（适合需要自定义命名的场景） #

如果不想创建多个Logger，你可以在模型内部的training_step、validation_step等方法里，给日志指标加上模型专属前缀，这样所有日志会存在同一个版本文件夹里，TensorBoard会自动按前缀分组。

示例（以Model1为例）： #

class Model1(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        # 添加前缀"Model1/"
        self.log("Model1/train_loss", loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        acc = accuracy(y_hat, y)
        self.log("Model1/val_acc", acc)

Model2同理，把前缀改成Model2/，然后用默认的Logger训练就行。这种方法的好处是日志都在同一个版本文件夹里，缺点是需要修改每个模型的日志代码。

两种方法都能实现你想要的同图对比效果，推荐第一种，更简洁且不需要改动模型逻辑～

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