LightingModel.save_checkpoint()

[Karesto]

Description & Motivation

Allow Lightning model to save a checkpoint, that you can load with load_from_checkpoint.

Pitch

I have encountered this issue several times, as I try to switch to Lightning, it is not simple to take old checkpoints from native pytorch and run them on the new code (for testing/comparison purposes for examples)
I tried to make a model, then load LightningModel.Model with the weights, but then if I want to save it, I need to run a full mock training which isn't handy or practical.

Alternatives

You could have a simple attach method on the trainer to attach a model without having to use Fit/Test/predict etc... and requiring arguments. Which would be a good alternative.

Additional context

No response

cc @lantiga @Borda

[SkafteNicki]

Hi @Karesto, thanks for opening this issue.
Regarding the opening of old checkpoints from native pytorch and loading them in a lightning module this is fully supported by using the load_state_dict method (native pytorch method):

pytorch_model = StandardPyTorchModel()
torch.save(pytorch_model.state_dict(), "checkpoint.pt")
lightning_model = LightningModel()
lightning_model.load_state_dict(torch.load("checkpoint.pt"))

If you want to manually save lightning compatible checkpoints without having to run a mock training something like this should work:

model = StandardPyTorchModel()
checkpoint = {
    "state_dict": model.state_dict(),
    "hyper_parameters": getattr(model, 'hparams', {}),
    "pytorch-lightning_version": pl.__version__,
    "epoch": 0,
    "global_step": 0,
}
torch.save(checkpoint, "checkpoint.ckpt")
loaded = SimpleModel.load_from_checkpoint("checkpoint.ckpt")