Hi,
Thanks for your question - ability to save/load checkpoints is an important feature, and it’s good for us to have some input on how people could be using it.
While we’re considering how to include this into Opacus API here’s what you need to know to make it work today.


PrivacyEngine doesn’t maintain links to model, optimizer, or data_loader. The only important state maintained by PrivacyEngine is accountant. Accountant’s state is just a list of numerical tuples, so torch.save() or any other pickle mechanism should do. That said, you probably should save the accountant (privacy_engine.accountant), not the privacy engine itself. That’s because we also maintain the link to the dataset used in the first call to the make_private() method - to do a sanity check the dataset is not being swapped in the middle (accounting is performed on per-dataset basis)
While GradSampleModule functionally is just a wrapper around nn.Module, saving/loading probably doesn’t work for them out of the box. I’d say your best bet is to save/load underlying model and then wrap it with GradSampleModule every time you’re restoring from the checkpoint.

To summarise, here are the steps you need to take
On saving:

Save accountant: torch.save(privacy_engine.accountant)

Save model: torch.save(model. _module.state_dict())

If your optimizer has state (e.g. dynamic lr or momentum) - save wrapped optimizer: optimizer.original_optimizer.state_dict()


On loading:

Initialize empty PrivacyEngine
Load accountant and replace brand new with the one you’ve just initialized: privacy_engine.accountant = accountant_you_have_just_loaded_from_checkpoint

Load your non-private nn.Module as normal
Load your non-private optimizer as normal
Pass loaded model and optimzier to privacy_engine.make_private()