Global step always zero after loading checkpoint

[timothybrooks]

I am saving checkpoints inside my module using self.trainer.save_checkpoint(path). I am able to load these checkpoints into the model using MyModel.load_from_checkpoint(path) and trainer using Trainer(resume_from_checkpoint=path). However, both the resulting model and trainer have global_step=0 regardless of the step when saving. From the documentation I was of the impression that checkpoints saved the global step, which is important for my use case. How can I attain the global step from a checkpoint?

[carmocca]

Here's a reproduction snippet

def test_bug(tmpdir):
    class TestModel(BoringModel):
        def on_pretrain_routine_end(self) -> None:
            # BUG
            # this fails because global_step is incorrectly increased
            assert self.trainer.global_step == 4

    model = BoringModel()
    trainer = Trainer(default_root_dir=tmpdir, max_epochs=2, limit_train_batches=2)
    trainer.fit(model)
    
    # correct
    # steps: (0, 1, 2, 3)
    # it is increased to 4 so it can run if fit is called again
    assert trainer.global_step == 4

    ckpt = tmpdir / "model.ckpt"
    trainer.save_checkpoint(ckpt)

    trainer = Trainer(default_root_dir=tmpdir, resume_from_checkpoint=str(ckpt), max_epochs=4, limit_train_batches=2)
    
    # correct
    # the checkpoint state will be restored at the start of fit
    # an argument could be made to change this
    assert trainer.global_step == 0

    trainer.fit(TestModel())
    
    # correct
    # steps: (4, 5, 6, 7)
    assert trainer.global_step == 8

Note that the global_step does get loaded, just after fit, not after the Trainer is instantiated (see the snippet)

However, this made me notice the step number is not correct.

This should probably be looked at with #5007. We need to clean the step/epoch management across runs

[timothybrooks]

I see, thanks for the explanation. I would argue that Trainer(resume_from_checkpoint=path).global_step and Model.load_from_checkpoint(path).global_step should reflect the true global step, rather than loading only during a call to trainer.fit(). Having the loaded global step at zero is inaccurate, and prevents the global step from being used in downstream processing and analysis that does not involve calling trainer.fit(). The Model.load_from_checkpoint(path).global_step case would be particularly valuable to me in two instances:

1. Post-processing of model that involves global step, such as taking weighted combinations of checkpoint network weights based on their global step.

2. Creating plots and visualizations which demonstrate change throughout training, such as computing expensive metrics (which are not run during training), analyzing feature activations, or exploring latent space manipulations, all as a function of the global step.

In the meantime, is there perhaps a workaround for directly loading the global step from the checkpoint?

[carmocca]

Yeah, what you are proposing is sound. I'll take a shot when I get the chance 👍

You can get the global_step like this: torch.load(ckpt)["global_step"] - 1 (-1 due to the bug mentioned above)