Loss not logged for 16bit precision and manual optimization

[timothybrooks]

When training a GAN and thus using manual optimization, I include these lines:

self.log("loss_d", loss_d, prog_bar=True)
self.log("loss_g", loss_g, prog_bar=True)

The logging works as intended, both for my logger (W&B) and to the command line progress bar so long as I use 32 bit precision. However, when I use 16 bit precision, no logging occurs to either. I have tried printing the losses, and they appear to be reasonable values and on the correct devices.

[timothybrooks]

I have also tried printing self._results and it appears to sometimes be missing entries when using 16 bit, whereas it always includes all entries in 32 bit. Also I disable 16 bit via with torch.cuda.amp.autocast(enabled=False): out of necessity in a couple select parts of the algorithm in case that is relevant.

[justusschock]

Hi @timothybrooks do you have an example to reproduce this behavior?

[timothybrooks]

Its entangled in a large project, but I've extracted a sketch of the training step which is the most relevant portion:

@property
def automatic_optimization(self) -> bool:
    return False

def training_step(...):
    
    # Discriminator.
    optimizer_idx = 0
    optimizer: optim.Adam = self.optimizers()[optimizer_idx]
    self.toggle_optimizer(optimizer, optimizer_idx)

    loss_d = self.compute_loss_d(...)
    self.log("loss_d", loss_d, prog_bar=True)

    optimizer.zero_grad()
    self.manual_backward(loss, optimizer)
    optimizer.step()
    self.untoggle_optimizer(optimizer_idx)

    # Generator.
    optimizer_idx = 1
    optimizer: optim.Adam = self.optimizers()[optimizer_idx]
    self.toggle_optimizer(optimizer, optimizer_idx)

    loss_g = self.compute_loss_g(...)
    self.log("loss_g", loss_g, prog_bar=True)

    optimizer.zero_grad()
    self.manual_backward(loss, optimizer)
    optimizer.step()
    self.untoggle_optimizer(optimizer_idx)

Portions of code inside compute_loss_d and compute_loss_g disable autocasting. I'm unsure how much of this is needed to reproduce the bug, and whether manual optimization is part of the issue or whether it is simply an issue with mixed precision.

[denadai2]

I have this bug as well. I do not find why but it's NOT related to manual optimization. It appears to be related to float=16

[tchaton]

Hey @denadai2 @timothybrooks,

I confirm I can reproduce the bug.

def test_multiple_optimizers_logging(tmpdir):
    """
    Tests that only training_step can be used
    """

    class TestModel(BoringModel):

        def __init__(self):
            super().__init__()
            self.automatic_optimization = False

        def training_step(self, batch, batch_idx):
            # Discriminator.
            optimizer_idx = 0
            optimizer = self.optimizers()[optimizer_idx]
            self.toggle_optimizer(optimizer, optimizer_idx)

            loss_d = self.loss(batch, self.layer(batch))
            self.log("loss_d", loss_d, prog_bar=True)

            optimizer.zero_grad()
            self.manual_backward(loss_d, optimizer)
            optimizer.step()
            self.untoggle_optimizer(optimizer_idx)

            # Generator.
            optimizer_idx = 1
            optimizer = self.optimizers()[optimizer_idx]
            self.toggle_optimizer(optimizer, optimizer_idx)

            loss_g = self.loss(batch, self.layer(batch))
            self.log("loss_g", loss_g, prog_bar=True)

            optimizer.zero_grad()
            self.manual_backward(loss_g, optimizer)
            optimizer.step()
            self.untoggle_optimizer(optimizer_idx)

            import pdb; pdb.set_trace()

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
            optimizer_2 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
            return optimizer, optimizer_2

    model = TestModel()
    model.training_epoch_end = None
    model.val_dataloader = None

    limit_train_batches = 2
    trainer = Trainer(
        default_root_dir=tmpdir,
        limit_train_batches=limit_train_batches,
        limit_val_batches=2,
        max_epochs=1,
        log_every_n_steps=1,
        weights_summary=None,
        gpus=1,
        precision=16,
    )

    trainer.fit(model)

    expected = {'epoch', 'loss_d', 'loss_g'}
    logged = set(trainer.logged_metrics.keys())
    assert expected == logged
    expected = {'loss_d', 'loss_g'}
    logged = set(trainer.progress_bar_metrics.keys())
    assert expected == logged