Correct way to save a trained model?

[carschandler]

I've looked around a decent bit and it's possible that I've somehow missed this, but what is the correct way to save a model which is performing well at the end of a training loop, and how would I go about re-using it later? I now understand how to save replay buffer contents, but to my understanding, this doesn't save any information about the weights of the models themselves, right?

I thought that maybe using torch.save might be the route to go, but what do I use it on? For example, I have a policy/actor which is a QValueActor with a DuelingCnnDQNet value net inside of it. I also have a policy_explore which is a Seq(policy, EGreedyModule). Is saving the QValueActor enough to get up and running again? Do I also need to save the policy_explore? What about the value net?

[jasperjonkhans]

Save the state_dict of the QValueActor. Because the DuelingCnnDQNet is registered as a submodule of the actor, its parameters are included recursively in the actor’s state_dict. You do not need to save the actor and its contained value network separately.

To load it, recreate the same QValueActor architecture and pass the saved state dictionary to load_state_dict(). For inference, those actor weights are normally sufficient, and you can reconstruct the exploration wrapper separately.

If you want to resume training rather than just run inference, also save the optimizer state, training step, random or exploration state, and any scheduler state. Saving the replay buffer is optional and is separate from saving the model weights.