How to Log Metrics (eg. Validation Loss) To TensorBoard when using Pytorch Lightning?

user_3023615 0
Viewed 1239

I am using Pytorch Lightning to train my models (on GPU devices, using DDP) and TensorBoard is the default logger used by Lightning.

My code is setup to log the training and validation loss on each training and validation step respectively.

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 correctly plots both the train_loss and val_loss charts in the SCALERS tab. However, in the HPARAMS tab, on the left side bar, only hp_metric is visible under Metrics.

enter image description here

However, in the HPARAMS tab, on the left side bar, only hp_metric is visible under Metrics.

enter image description here

How can we add train_loss and val_loss to the Metrics section? This way, we will be able to use val_loss in the PARALLEL COORDINATES VIEW instead of hp_metric.

Image showing hp_metric and no val_loss: enter image description here

Using Pytorch 1.8.1, Pytorch Lightning 1.2.6, TensorBoard 2.4.1

pytorch-lightning pytorch tensorboard machine-learning python
1 Answers
  • You can use self.logger.log_hyperparams method to log hyperparameters and metrics in tensorboard. (see pytorch lightning tensorboard docs)
  • The values you added by self.log will be displayed in hparam plugin if and only if you have passed same key names in metric. (see pytorch tensorboard docs)

Sample code (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)
user_176703800
Related