python - 使用 Pytorch Lightning 时如何将指标(例如验证损失)记录到 TensorBoard？

Question

我正在使用 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 选项卡的左侧栏中，只有 hp_metric 在 Metrics 下可见。

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

我们如何将 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

Answer 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)