How to save density _estimator after the NN is trained?

[jsv1206]

Hello!

Thanks for the amazing package and the documentation.

I am using SNRE to get the posterior and looking into different samplers following this example.

https://sbi-dev.github.io/sbi/tutorial/11_sampler_interface/

I would like to save the density_estimator after training because I want to use the same trained network to get potential_fn and posterior for different sampling algorithms.

However, I am not sure how to do this given that the neural network class is built and hidden in /sbi/sbi/neural_nets/classifier.py

This is the code I have with an embedding network for my inputs

neural_classifier = utils.classifier_nn(model='resnet', 
                                      embedding_net_x=embedding_net,
                                      hidden_features=250,
                                      num_blocks=5,
                                       )
inference = SNRE(prior, classifier=neural_classifier, device='cuda:0')
density_estimator = inference.append_simulations(theta, input_features).train()

Is there a way to do something like:

# save the trained network
torch.save(density_estimator.state_dict(), 'snre_test_save_model.pt')

# load this to compute further steps
model.load_state_dict(torch.load('snre_test_save_model.pt'), strict=False)
model.eval()

Thanks!

[michaeldeistler]

Hi there, it should be possible to pickle the density estimator. I.e.

import pickle

# Write.
with open("path.pkl", "wb") as handle:
    pickle.dump(density_estimator, handle)

# Read.
with open("path.pkl", "rb") as handle:
    density_estimator = pickle.load(handle)

(tested this on CPU only, please let me know if it leads to issues for you)

[jsv1206]

Thanks a lot, Michael!
This seems to work.