Pytorch-Lightning save and continue training from state_dict.
🚀 Feature
Save the model and other checkpoint at any step as a dict. And load these points to retraining the model on datasets.
Motivation
Recently, I am working on federated learning, an learning paradiam where the training is runing on different clients. And the clients share one common model which is aggregate from local client models at server.  This training precedure ask the local clients could stop and send the middle models to the server after a given epoch or steps end. And then these middle models would be aggregated at server to get the shared common model. Next, the clients load common model and continue training. This process would execute several rounds.
Pitch
For a LightningModule, I can get its state_dict at anytime. And the Trainer could suspend its traning procedure and recover traning from state_dict sended from server. I tested trainer.fit()  by calling two times, with max_epochs. The model could retraining from a given state_dict. But for max_steps, the trainer will execute the second trainer.fit()  from the first batch and finish after one epoch. Not from the  last breakpoint to next max_steps.