使用 Pytorch Lightning 时如何将指标（例如验证损失）记录到 TensorBoard？答案

【问题标题】：How to Log Metrics (eg. Validation Loss) To TensorBoard when using Pytorch Lightning?使用 Pytorch Lightning 时如何将指标（例如验证损失）记录到 TensorBoard？
【发布时间】：2021-06-30 20:48:26
【问题描述】：

我正在使用 Pytorch Lightning 来训练我的模型（在 GPU 设备上，使用 DDP），TensorBoard 是 Lightning 使用的默认记录器。

我的代码设置为分别记录每个训练和验证步骤的训练和验证损失。

class MyLightningModel(pl.LightningModule):

    def training_step(self, batch):
        x, labels = batch
        out = self(x)
        loss = F.mse_loss(out, labels)
        self.log("train_loss", loss)
        return loss

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

TensorBoard 在SCALERS 选项卡中正确绘制了train_loss 和val_loss 图表。但是，在 HPARAMS 选项卡的左侧栏上，Metrics 下只有 hp_metric 可见。

但是，在 HPARAMS 选项卡的左侧栏中，Metrics 下只有 hp_metric 可见。

我们如何将train_loss 和val_loss 添加到Metrics 部分？这样，我们就可以在PARALLEL COORDINATES VIEW 中使用val_loss 而不是hp_metric。

图片显示hp_metric，没有val_loss：

使用 Pytorch 1.8.1、Pytorch Lightning 1.2.6、TensorBoard 2.4.1

【问题讨论】：

标签： python machine-learning pytorch tensorboard pytorch-lightning

【解决方案1】：

您可以使用self.logger.log_hyperparams 方法在 tensorboard 中记录超参数和指标。（见pytorch lightning tensorboard docs）
当且仅当您在metric 中传递了相同的键名时，self.log 添加的值才会显示在 hparam 插件中。（见pytorch tensorboard docs）

示例代码 (full code)：

class BasicModule(LightningModule):
    def __init__(self, lr=0.01):
        super().__init__()
        self.model = models.resnet18(pretrained=False)
        self.criterion = nn.CrossEntropyLoss()
        self.lr = lr
        self.save_hyperparameters()
        
        metric = MetricCollection({'top@1': Accuracy(top_k=1), 'top@5': Accuracy(top_k=5)})
        self.train_metric = metric.clone(prefix='train/')
        self.valid_metric = metric.clone(prefix='valid/')
    
    def on_train_start(self) -> None:
        # log hyperparams
        self.logger.log_hyperparams(self.hparams, {'train/top@1': 0, 'train/top@5': 0, 'valid/top@1': 0, 'valid/top@5': 0})
        return super().on_train_start()
    
    def training_step(self, batch, batch_idx, optimizer_idx=None):
        return self.shared_step(*batch, self.train_metric)

    def validation_step(self, batch, batch_idx):
        return self.shared_step(*batch, self.valid_metric)

    def shared_step(self, x, y, metric):
        y_hat = self.model(x)
        loss = self.criterion(y_hat, y)
        self.log_dict(metric(y_hat, y), prog_bar=True)
        return loss

if __name__ == '__main__':
    # default_hp_metric=False
    logger = loggers.TensorBoardLogger('', 'lightning_logs', default_hp_metric=False)
    trainer = Trainer(max_epochs=2, gpus='0,', logger=logger, precision=16)

【讨论】：