SBI with high-dimensional simulations with TSNPE?

[Millet21]

Hi!,

I'm working on a project where I need to fit a model to high-dimensional data (specifically, video recordings with 1000 frames). The model has between 6 and 10 parameters to infer, and I've chosen to use SNPE for this task.

However, I'm running into a couple of challenges:

Due to resource constraints, I can’t generate more than ~10,000 simulations.

Simulation and training (including both the neural density estimator and the embedding network) at the same time are quite computationally expensive on my GPU.

Has anyone faced similar issues, or do you have suggestions for how to make SNPE more efficient in this setting? Any advice, ideas, tips, or other tricks would be greatly appreciated :)

Thanks in advance!

[michaeldeistler]

Hi there,

thanks for reaching out! I don't think I have concrete tips for you though --- video data is just expected to make training expensive.

For maximal flexibility over the dataloader and training loop, I recommend to check out this tutorial.

Aside: 10k simulations is really on the lower end of what's necessary to make SNPE work reliably, especially with a large embedding network. I don't think it's impossible, but you might have to do quite some hyperparameter tuning etc.

Hope this helps!
Michael

[Millet21]

Hi,
Thanks for your response! :) You're right—it turned out to be more challenging than I expected, especially with such large models.
I've tried reducing the model size and limiting the number of parameters (currently down to six), but I'm still running into issues with SNPE. Specifically, I keep running out of GPU memory during the second round of training :(
Are there any examples of efficient GPU management within an SNPE loop that uses an embedding network that I could look at for inspiration?
Thanks in advance,
Mathis

[michaeldeistler]

You can probably overcome issues with memory by using TSNPE:

from sbi.inference import NPE
from sbi.utils import RestrictedPrior, get_density_thresholder

inference = NPE(prior)
proposal = prior
for _ in range(num_rounds):
    theta = proposal.sample((num_sims,))
    x = simulator(theta)
    _ = inference.append_simulations(theta, x).train(force_first_round_loss=True)
    posterior = inference.build_posterior().set_default_x(x_o)

    accept_reject_fn = get_density_thresholder(posterior, quantile=1e-4)
    proposal = RestrictedPrior(prior, accept_reject_fn, sample_with="rejection")

Aside, TSNPE is typically much more stable and robust than other SNPE methods, and it is also the only SNPE method for which you can easily implement your own training loop (following this tutorial, where you just have to add the loop over rounds).

As for GPU data management, please read this. The key take-away is to pass device="cuda" (or "gpu") to NPE(...) and to pass data_device="cpu" to append_simulations.

I hope this helps, but let me know if there are any issues!
Michael

[Millet21]

Thank you very much :) I will try this !
Mathis

[Millet21]

Hey,

Sorry to disturb you again. The T-SNPE procedure is working pretty well, but I still have some errors i don't understand ... :/ (cf. screenshots below)

Even though I tried to do it in batches, sampling the restricted distribution is sometimes consuming a lot of memory.

Can you help me with this?

Thanks in advance!
Mathis

[michaeldeistler]

Hi!

I recommend to not use simulate_for_sbi here, in order to get a bit more flexibility. Instead of it, do:

# The default for `max_sampling_batch_size ` is 10_000, which (I guess) is causing your issues.
theta = proposal.sample((10_000,), max_sampling_batch_size=100)

# Run the simulations.
simulation_outputs: list[Tensor] = []
batches = torch.split(theta, simulation_batch_size)
for batch in tqdm(batches, disable=not show_progress_bar):
    simulation_outputs.append(simulator(batch))
x = torch.cat(simulation_outputs, dim=0)
theta = torch.as_tensor(theta, dtype=float32)

Let me know if this helps!
Michael

[Millet21]

Yes, that was exactly what I was looking for!
Thanks for all

[Millet21]

Hi all,
I have a few questions regarding the Truncated SNPE procedure : is it possible to run SBC and compute the MAP in this case? Or is it just not implemented yet but completely feasible in practice?
Thanks!

[michaeldeistler]

Hi!

SBC is indeed not implemented (well, it is implemented, but only for the first round). To implement it, you need a method that sets posterior.log_prob().exp() to zero outside of the support of the RestrictedPrior, and then run SBC with parameters generated from the RestrictedPrior.

MAP should be implemented, are you experiencing issues with it?

Michael

[Millet21]

Ok, I will try to do so :)
And yes, MAP is working well, my bad!
Mathis