Show how to over-fit batches for real about lightning

[aakashvardhan]

I think you can call the trainer class property 'logged_metrics', which calls the train loss

for example:

# Retrieve logged losses
    train_losses = trainer.logged_metrics['train_loss']
    
    # Check if the final loss is lower than the initial loss
    initial_loss = train_losses[0]
    final_loss = train_losses[-1]
    
    assert final_loss < initial_loss,

[svnv-svsv-jm]

Not really. logged_metrics only holds the last value:

print(logged_metrics)
>>> {'train_loss': tensor(0.0162), 'train_loss_step': tensor(0.0162), 'train_loss_epoch': tensor(0.0162)}

So I had to do something like:

# Run a fake training just to get the loss over one batch at the start (`max_steps=1`)
    trainer = Trainer(
        logger=False,
        enable_checkpointing=False,
        accelerator="auto",
        overfit_batches=1,
        max_steps=1,
    )
    trainer.fit(model, loader)
    logged_metrics = trainer.callback_metrics
    logger.info(logged_metrics)
    loss_before = logged_metrics["train_loss"]

    # Now overfit on the same batch, but for more steps and get the loss
    trainer = Trainer(
        logger=False,
        enable_checkpointing=False,
        accelerator="auto",
        overfit_batches=1,
        **kwargs,
    )
    trainer.fit(model, loader)
    logged_metrics = trainer.callback_metrics
    logger.info(logged_metrics)
    loss_after = logged_metrics["train_loss"]

which is not nice.

Either an example of how this is intended to be done should be provided, or handle specially the case when overfit_batches>0.