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!

[ali-akhavan89]

Hello,

Have the features discussed here already been implemented into SBI?

I am using pairs of ground truth parameters and time series data to train a density estimator in an amortized way (NPE). For some complex and high-dimensional models, I need a very large number of datasets to achieve proper training. Is it possible to resume the training when the previous round has converged and, at the same time, remove the old datasets and introduce a new batch without losing or restarting the network parameters?

If this is possible, I have a follow-up clarification question. What happens to the learning rate in between the training rounds? Is there a way to control the learning rate in a way that it is not restarted each round a new dataset is introduced?

Thank you so much for your support!
Ali

[michaeldeistler]

Hi there,

In general, if you need this level of control over the exact training loop, I would recommend you to write the training loop yourself, which we explain here.

If you want to stick to the sbi pre-configured training loop:

• for removing data, the answer is still the same as above, i.e.: Removing data is not easily supported, but you could remove the data from the ._theta_roundwise, ._x_roundwise, and ._prior_masks_roundwise, see here.
• if train(..., retrain_from_scratch=False), then you will not lose/restart the network parameters
• sbi does not schedule the learning rate at all. Therefore there is no such thing as "learning rate restarting".

Hope this helps!
Michael