How can be the rng_state stored using the checkpoint_callbak.

[dhorka]

Hi,

I would like to know how can I store the rng_state from the different random generators used in a model, such as NumPy and PyTorch random generators. Each framework has its own method to recover the states, for instance, PyTorch has torch.get_rng_state that returns the current state of the rng. I would like to be able to save these states, in order to resume the experiment at the same exact point where it was at the first run. As far as I saw, at this moment this behaviour is not implemented in the checkpoint_callback.

Thanks,

[SkafteNicki]

Lightnings checkpoints also save the users hyperparameters (https://pytorch-lightning.readthedocs.io/en/latest/common/weights_loading.html). Would that be an option for you to save the states as part of the hyperparameters?

class MyModule(LightningModule):
    def __init__(self, torch_state, numpy_state):
        self.save_hyperparameters()

numpy_state = numpy.random.get_state()
torch_state = torch.get_rng_state()

model = MyModule(torch_state, numpy_state)

[dhorka]

Hmmm, I have to think about it. I mean, the rng state should be obtained at the same moment as the checkpoint is saved. This is because I need to have the last state to restart the experiment at the same point. Moreover, I should load this rng state at each process in the case of multi-GPU ddp. I suppose I can try to save the rng states as hyperparameters at the end of val_epoch (in case I am generating a checkpoint at the end of val_epoch). Then my doubt is where and how I should set the rng_state to be able to recover the experiment in the same exact point.

In vanilla pytorch, I have a function that encapsulates these things and I call it when I want to save a checkpoint. When loading, I load the checkpoint and then recover the model, optimizers states, etc. and just before to start the training loop, I set the rng states. To avoid that, any random method that can be called moves forward my rng_state. By doing this, I can resume the experiment at the same point as it was before stopping the experiment. In pytorch lightning I am not sure how to do this saving the states as a hyperparameter.

[SkafteNicki]

What about writing your own callback and using on_save_checkpoint method to save the states and then load them in your model using the on_load_checkpoint method:

class StateSaver(Callback):
    def on_save_checkpoint(trainer, pl_module, checkpoint):
        checkpoint['torch_state'] = torch.get_rng_state()
        checkpoint['numpy_state'] = numpy.random.get_state()

and then in your model

class MyModel(LightningModule):
    def on_load_checkpoint(checkpoint):
        self.torch_state = checkpoint['torch_state'].pop()
        self.numpy_state = checkpoint['numpy_state'].pop()

[dhorka]

Oh! I did not know about these methods. I think that it can solve the issue. Just a couple of doubts. Why the first one is a callback, and the second is just inside the model? Which is the difference between on_load_checkpoint and on_pretrain_routine_start? I mean, I suppose that on_load_checkpoint is called before on_pretrain_routine_start, isn't it? I would like to call the torch.set_rng_state(self.torch_state) just before the training loop starts. Is this possible? And the final question, both calls the on_load and on_save are only executed on rank=0 in multi-gpu scenarios, isn't? or should I control this scenario?

[dhorka]

Answering my own questions:

I would like to call the torch.set_rng_state(self.torch_state) just before the training loop starts. Is this possible?
Yes, is possible. I just need to put the torch.set_rng_state in the hook on_fit_start()

Related to on_pretrain_routine_start, I just found that is executed after the on_fit_start() hook. However, I still do not know when is useful to use this hook.

And the final question, both calls the on_load and on_save are only executed on rank=0 in multi-gpu scenarios, isn't?

It is not I put a print sentence on the on_save_checkpoint and the hook is called on all the ranks... Then I suppose I should control it. Isn't it? I suppose that is the expected behavior, right? I mean I suppose if in some case you want to store some specific information of each rank you can do it.

I still do not understand why you did the first one as a callback. In fact, I do not understand when I should use the callbacks.

Thanks,