NeuralPosteriorEnsemble & multi-rounds

[jecampagne]

Hello,
Is it possible to use NeuralPosteriorEnsemble as in the snippet

theta, x = simulate_for_sbi(cond_simulator, prior, num_simulations=50_000)
def train_nn(rng):
    utils.sbiutils.seed_all_backends(rng)
    density_estimator_build_fun = posterior_nn(model="maf",  hidden_features=64)
    inference = SNPE(prior=prior, device=device,
                     density_estimator=density_estimator_build_fun)
    _ = inference.append_simulations(theta, x).train()
    return inference.build_posterior()

posteriors = Parallel(n_jobs=-1)(delayed(train_nn)(rng) for rng in random_state)
final_posterior = NeuralPosteriorEnsemble(posteriors)

but with a multi-rounds for each trained posterior?

[janfb]

Hi,
yes, just add a loop over the rounds as in
https://github.com/mackelab/sbi/blob/main/tutorials/03_multiround_inference.ipynb

[jecampagne]

Hi, yes, just add a loop over the rounds

You mean that the new train_nn function should look like

def train_nn(rng, num_rounds, num_simul):
    utils.sbiutils.seed_all_backends(rng)

    density_estimator_build_fun = posterior_nn(model="maf",  hidden_features=64)
    inference = SNPE(prior=prior, device=device,
                     density_estimator=density_estimator_build_fun)

   proposal = prior
   for _ in range(num_rounds):
      theta, x = simulate_for_sbi(simulator, proposal, num_simul)
      # only for SNLE & SNPE proposal=proposal
      density_estimator = inference.append_simulations(
            theta, x, proposal=proposal
        ).train()
       posterior = inference.build_posterior(density_estimator)
       proposal = posterior.set_default_x(x_o)

    return posterior

[janfb]

yes.
But you should make sure you are using the same x_o throughout all calls.
and beware that you are simulating num_simul in each round.

[jecampagne]

yes. But you should make sure you are using the same x_o throughout all calls.

Yes, that's the case

and beware that you are simulating num_simul in each round.

Well what do you mean?

[janfb]

the total number of simulations will be num_rounds * num_simul and not just num_simulations as in your first snippet.

[jecampagne]

Well I have a crash on Google Colab (CPU):

def train_nn(rng, num_rounds, num_simul):
    utils.sbiutils.seed_all_backends(rng)

    density_estimator_build_fun = posterior_nn(model="maf",  hidden_features=64)
    inference = SNPE(prior=prior, device=device,
                     density_estimator=density_estimator_build_fun)

    proposal = prior
    for _ in range(num_rounds):
      theta, x = simulate_for_sbi(cond_simulator, proposal, num_simul)
      # only for SNLE & SNPE proposal=proposal
      density_estimator = inference.append_simulations(
            theta, x, proposal=proposal
        ).train()
      posterior = inference.build_posterior(density_estimator)
      proposal = posterior.set_default_x(x_o)

    return posterior

ensemble_size = 3
np.random.seed(10)
random_state = np.random.randint(1_000_000, size=ensemble_size)

num_rounds=2
num_simul=10_000
posteriors = Parallel(n_jobs=-1)(delayed(train_nn)(rng,num_rounds,num_simul) for rng in random_state)
final_posterior = NeuralPosteriorEnsemble(posteriors)

TypeError                                 Traceback (most recent call last)

[<ipython-input-35-1d6ed27e5fca>](https://localhost:8080/#) in <module>
      1 num_rounds=2
      2 num_simul=10_000
----> 3 posteriors = Parallel(n_jobs=-1)(delayed(train_nn)(rng,num_rounds,num_simul) for rng in random_state)
      4 final_posterior = NeuralPosteriorEnsemble(posteriors)

4 frames

[/usr/lib/python3.7/concurrent/futures/_base.py](https://localhost:8080/#) in __get_result(self)
    382     def __get_result(self):
    383         if self._exception:
--> 384             raise self._exception
    385         else:
    386             return self._result

TypeError: len() of unsized object

After investigation it seems that the statement that crashes is proposal = posterior.set_default_x(x_o) while x_o is well defined.

[janfb]

can you show the full traceback? where is len() called?

[jecampagne]

I have exectuted the train_nn function w/o the joblib Parallel and it fails also

def train_nn(rng, num_rounds, num_simul):
    print("train, rng: ",rng)
    utils.sbiutils.seed_all_backends(rng)


    density_estimator_build_fun = posterior_nn(model="maf",  hidden_features=64)
    inference = SNPE(prior=prior, device=device,
                     density_estimator=density_estimator_build_fun)

    proposal = prior
    for ir in range(num_rounds):
      print("train, rng: ",rng, " round:", ir)
      theta, x = simulate_for_sbi(cond_simulator, proposal, num_simul)
      # only for SNLE & SNPE proposal=proposal
      density_estimator = inference.append_simulations(
            theta, x, proposal=proposal
        ).train()
      print("\n build_posterior")
      posterior = inference.build_posterior(density_estimator)
      print("new proposal")
      proposal = posterior.set_default_x(x_o)
      print("end this round")

    print("all done..")
    return posterior


tmp = train_nn(random_state[0], num_rounds=1, num_simul=100)

leads to

train, rng:  345353
train, rng:  345353  round: 0

Running 100 simulations.: 100%
100/100 [00:00<00:00, 1328.91it/s]

 Neural network successfully converged after 134 epochs.
 build_posterior
new proposal

---------------------------------------------------------------------------

IndexError                                Traceback (most recent call last)

[/usr/local/lib/python3.7/dist-packages/jax/_src/device_array.py](https://localhost:8080/#) in __len__(self)
    241     try:
--> 242       return self.aval.shape[0]
    243     except IndexError as err:

IndexError: tuple index out of range


The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)

4 frames

[<ipython-input-45-c0afa56b14f8>](https://localhost:8080/#) in <module>
----> 1 tmp = train_nn(random_state[0], num_rounds=1, num_simul=100)

[<ipython-input-44-144ee18e600b>](https://localhost:8080/#) in train_nn(rng, num_rounds, num_simul)
     19       posterior = inference.build_posterior(density_estimator)
     20       print("new proposal")
---> 21       proposal = posterior.set_default_x(x_o)
     22       print("end this round")
     23 

[/usr/local/lib/python3.7/dist-packages/sbi/inference/posteriors/base_posterior.py](https://localhost:8080/#) in set_default_x(self, x)
    126         """
    127         self._x = process_x(
--> 128             x, x_shape=self._x_shape, allow_iid_x=self.potential_fn.allow_iid_x
    129         ).to(self._device)
    130         self._map = None

[/usr/local/lib/python3.7/dist-packages/sbi/utils/user_input_checks.py](https://localhost:8080/#) in process_x(x, x_shape, allow_iid_x)
    568     """
    569 
--> 570     x = atleast_2d(torch.as_tensor(x, dtype=float32))
    571 
    572     input_x_shape = x.shape

[/usr/local/lib/python3.7/dist-packages/jax/_src/device_array.py](https://localhost:8080/#) in __len__(self)
    242       return self.aval.shape[0]
    243     except IndexError as err:
--> 244       raise TypeError("len() of unsized object") from err  # same as numpy error
    245 
    246   setattr(device_array, "__len__", __len__)

TypeError: len() of unsized object

[jecampagne]

Ha! I found the problem... it is my fault (of course) : using a JAX simulator, I must use

posterior.set_default_x(x_o_tensor)

with

x_o_tensor = torch.tensor(np.asarray(x_o), device=device)[None]

Now it is working :)