Batch learning

[Kojobu]

Hi all,
I'm trying to implement a sbi using an embedding_net. Due to the high-dim of my data the following steps are necessary:

1. Precompute the prior samples and run the corresponding likelihood samples on a cluster.
2. Pretrain both the summary_net and 2.5) pretrain the density estimation net.
3. Train both together

2.5) and 3) causing issues with this package.

2.5) How can I save and load the weights of the density estimation network?

3. I'm training both nets together via inference's append_simulations method which soon runs out of memory due to the size of my likelihood samples. Is there a way to delete the samples saved in the inference objects? Or even better: Can I pass my own dataloader somehow?

Help is very much appreciated,
Have a wonderful day
~ Tom S

[michaeldeistler]

Hi there,

thanks for reaching out!

2.5.: See here for how to save and load posteriors.
3. If you are on GPU, then do append_simulations(..., data_device="cpu"), see here. Removing data is not easily supported, but you could remove the data from the ._theta_roundwise, ._x_roundwise, and ._prior_masks_roundwise, see here. Writing your own training loop is possible only if you are working on the github version of sbi. This is not released on pypi yet. Here is a tutorial on how to do that.

Michael

[Kojobu]

Especially the hint how to write your own training loop helped me tremendously! Thank you for the quick reply and thanks for maintaining this package!

[Kojobu]

Just another quick question: Is there a way to save the weights of the density estimator or more precisely, can I concatenate the embedding_net after I build the posterior_nn? I want to pretrain both the embedding_net and the density estimator and then train both together due to the complexity of my problem. Is there a way of doing so?

[michaeldeistler]

There is no easy way to do this, you will have to hack this in somehow.

[Kojobu]

Ok I got it to work now with deleting the data out of the inference class. I also tried writing my own trainingsloop but I ran into serious problems and due to the lack of documentation of this version I felt lost soon (I was using an embedding-net and couldn't get the dimensions of the conditions working). If I may give a recommendation for future versions:
It would help tremendously if one could pass a custom dataloader to the inference. For my case, my database is far too large to even put it into the ram. So I'm forced to use force_first_round_loss=True + deleting the data every batch.
It works for now but it's not very elegant and rather slow.

[janfb]

Thanks for the suggestion with the dataloader, that's good to know. Can you please create an issue / feature request for this?
Then we can work on it for the next release. Thanks!