使用 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在选项卡中正确绘制train_loss和图表。但是，在选项卡中的左侧栏上，仅在 下可见。val_lossSCALERSHPARAMShp_metricMetrics

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

我们如何将train_loss和添加val_loss到该Metrics部分？这样，我们就可以使用val_lossinPARALLEL COORDINATES VIEW来代替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方法在张量板中记录超参数和指标。（参见pytorch 闪电张量板文档）
• self.log当且仅当您在 中传递了相同的键名时，您添加的值才会显示在 hparam 插件中metric。（参见pytorch 张量板文档）

示例代码（完整代码）：

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)