Logging tensorboard not showing loss

[rienboonstoppel]

I am trying to get my losses in Tensorboard, but I am quite confused.
I simply return a dict after training_step and validation_step, containing loss and log. However, the only thing that is showing up in Tensorboard of that run is the 'hp_metric' thing... Nothing on the scalars of the losses... Both during and after training.
I got it working with manual logging (self.log(...)), however, that should not be necessary right? And is more complicated when I want for example training and validation loss in one plot. I am working on a Super Resolution GAN. This is my trainer:

class LitTrainer(pl.LightningModule):
    def __init__(self,
                 netG,
                 netD,
                 lr: float = 0.0002,
                 b1: float = 0.5,
                 b2: float = 0.999,
                 **kwargs
                 ):
        super().__init__()
        self.save_hyperparameters(ignore=["netG", "netD"])

        self.netG = netG
        self.netD = netD

        self.criterion_GAN = GANLoss("vanilla")
        self.criterion_edge = edge_loss1
        self.criterion_pixel = torch.nn.L1Loss()

    def forward(self, inputs):
        return self.netG(inputs)

    def prepare_batch(self, batch):
        return batch["LR"][tio.DATA].squeeze(4), batch["HR"][tio.DATA].squeeze(4)

    def training_step(self, batch, batch_idx, optimizer_idx):
        imgs_lr, imgs_hr = self.prepare_batch(batch)

        # train generator
        if optimizer_idx == 0:
            self.gen_hr = self(imgs_lr)

            loss_adv = self.criterion_GAN(self.netD(self.gen_hr), True)
            loss_edge = self.criterion_edge(self.gen_hr, imgs_hr)
            loss_pixel = self.criterion_pixel(self.gen_hr, imgs_hr)
            g_loss = loss_adv + loss_edge + loss_pixel

            # self.log("loss/G train", g_loss, on_step=True, on_epoch=True)
            tensorboard_logs = {"loss_g": {"train": g_loss}}
            return {"loss": g_loss, "log": tensorboard_logs}

        # train discriminator
        if optimizer_idx == 1:

            # for real image
            pred_real = self.netD(imgs_hr)
            real_loss = self.criterion_GAN(pred_real, True)
            # for fake image
            pred_fake = self.netD(self.gen_hr.detach())
            fake_loss = self.criterion_GAN(pred_fake, False)

            d_loss = (real_loss + fake_loss) / 2
            tensorboard_logs = {"loss_d": {"train": d_loss}}

            # self.log("loss/D train", d_loss, on_step=True)
            return {"loss": d_loss, "log": tensorboard_logs}

    def validation_step(self, batch, batch_idx):
        with torch.no_grad():
            imgs_lr, imgs_hr = self.prepare_batch(batch)
            gen_hr = self(imgs_lr)
            loss_adv = self.criterion_GAN(self.netD(gen_hr), True)
            loss_edge = self.criterion_edge(gen_hr, imgs_hr)
            loss_pixel = self.criterion_pixel(gen_hr, imgs_hr)
            g_loss = loss_adv + loss_edge + loss_pixel

            # for real image
            pred_real = self.netD(imgs_hr)
            real_loss = self.criterion_GAN(pred_real, True)
            # for fake image
            pred_fake = self.netD(self.gen_hr.detach())
            fake_loss = self.criterion_GAN(pred_fake, False)

            d_loss = (real_loss + fake_loss) / 2
            tensorboard_logs = {"loss_g": {"val": g_loss},
                                "loss_d": {"val": d_loss},
                                }

            # self.log("loss/G val", g_loss, on_step=True)
            # self.log("loss/D val", d_loss, on_step=True)

            return {"log": tensorboard_logs}

    def configure_optimizers(self):
        lr = self.hparams.lr
        b1 = self.hparams.b1
        b2 = self.hparams.b2
        opt_g = torch.optim.Adam(self.netG.parameters(), lr=lr, betas=(b1, b2))
        opt_d = torch.optim.Adam(self.netD.parameters(), lr=lr, betas=(b1, b2))
        return [opt_g, opt_d], []

And this is how I start training:

logger = TensorBoardLogger("log", name="test")

model = LitTrainer(netG=generator, netD=discriminator)
trainer = pl.Trainer(gpus=1, max_epochs=1, logger=logger, log_every_n_steps=10)
trainer.fit(model, train_dataloaders=training_loader, val_dataloaders=val_loader)

[awaelchli]

Are you using the latest Lightning version? If yes, this won't work. The new way for logging is through self.log instead of returning it from the step methods. See the docs here.

[rienboonstoppel]

Ah okay, then I was probably just looking at old examples... My bad!
But do you know how to log training and validation loss to the same graph in Tensorboard (preferably per epoch, on step is not necessary)? Because when I use self.log(...) in the seperate train_step and validation_step I get for every scalar a different graph. The solution I found online for the old version was to return nested dicts in the different step methods, however, that is not possible anymore if I am correct?

[tchaton]

You could use self.log("loss/train") and self.log("loss/val")

[rienboonstoppel]

No, I tried... That only places multiple figures under the same sub-header in Tensorboard.

I want the train loss and val loss to plot in the same figure. In barebones Pytorch with tensorboard you need to log them in a dictionary like this:

            self.writer.add_scalars(
                'Generator epoch loss', {'Train': train_epoch_loss_G,
                                         'Val': val_epoch_loss_G,
                                         }, epoch)
            self.writer.add_scalars(
                'Discriminator epoch loss', {'Train': train_epoch_loss_D,
                                             'Val': val_epoch_loss_D,
                                             }, epoch)

Is there a way in Lightning to do this?

[tchaton]

Could you try this self.log("loss", {"train": loss}) and self.log("loss", {"val": loss})

[rienboonstoppel]

Yes thanks!! That does the job!