Repository: sbi-dev/sbi

Issue:
Paralell simulation with simulatoe_for_sbi is not seeded
Hi, 
Looking at the `simulate_for_sbi`  I was wandering when `simulate_in_batches`  is used  with num_workers>1 how the seeds are managed ?
Thanks


Discussion:
Hi, 
at the moment we do not seed simulations generated with `simulate_for_sbi` internally. It's true that when using `num_workers>1` it is difficult to seed simulations from the outside. 
If this is essential for you, one workaround would be to pass seeds as extra parameters to your simulator and seed manually in your simulator. 

In the long term it would be nice to add the seeding feature to the `num_workers>1` scenario in `sbi`. 

Resolving patch:
diff --git a/sbi/simulators/simutils.py b/sbi/simulators/simutils.py
--- a/sbi/simulators/simutils.py
+++ b/sbi/simulators/simutils.py
@@ -3,7 +3,7 @@
 
 
 import contextlib
-from typing import Callable
+from typing import Callable, Optional
 
 import joblib
 import torch
@@ -11,12 +11,15 @@
 from torch import Tensor
 from tqdm.auto import tqdm
 
+from sbi.utils.sbiutils import seed_all_backends
+
 
 def simulate_in_batches(
     simulator: Callable,
     theta: Tensor,
     sim_batch_size: int = 1,
     num_workers: int = 1,
+    seed: Optional[int] = None,
     show_progress_bars: bool = True,
 ) -> Tensor:
     r"""
@@ -32,6 +35,7 @@ def simulate_in_batches(
             the entire theta in a single batch. When using multiple workers, increasing
             this batch size can further speed up simulations by reducing overhead.
         num_workers: Number of workers for multiprocessing.
+        seed: seed for reproducibility.
         show_progress_bars: Whether to show a progress bar during simulation.
 
     Returns:
@@ -39,6 +43,7 @@ def simulate_in_batches(
     """
 
     num_sims, *_ = theta.shape
+    seed_all_backends(seed)
 
     if num_sims == 0:
         x = torch.tensor([])
@@ -48,6 +53,13 @@ def simulate_in_batches(
         batches = torch.split(theta, sim_batch_size, dim=0)
 
         if num_workers != 1:
+            batch_seeds = torch.randint(high=1_000_000, size=(len(batches),))
+
+            # define seeded simulator.
+            def simulator_seeded(theta: Tensor, seed) -> Tensor:
+                seed_all_backends(seed)
+                return simulator(theta)
+
             # Parallelize the sequence of batches across workers.
             # We use the solution proposed here: https://stackoverflow.com/a/61689175
             # to update the pbar only after the workers finished a task.
@@ -58,9 +70,10 @@ def simulate_in_batches(
                     desc=f"Running {num_sims} simulations in {len(batches)} batches.",
                     total=len(batches),
                 )
-            ) as progress_bar:
+            ) as _:
                 simulation_outputs = Parallel(n_jobs=num_workers)(
-                    delayed(simulator)(batch) for batch in batches
+                    delayed(simulator_seeded)(batch, batch_seed)
+                    for batch, batch_seed in zip(batches, batch_seeds)
                 )
         else:
             pbar = tqdm(
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -2,16 +2,15 @@
 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
 import logging
+import random
 import warnings
 from math import pi
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type, Union
 
-import arviz as az
+import numpy as np
 import pyknos.nflows.transforms as transforms
 import torch
 import torch.distributions.transforms as torch_tf
-from arviz.data import InferenceData
-from numpy import ndarray
 from pyro.distributions import Empirical
 from torch import Tensor
 from torch import nn as nn
@@ -855,3 +854,17 @@ def gradient_ascent(
         return argmax_, max_val  # type: ignore
 
     return theta_transform.inv(best_theta_overall), max_val  # type: ignore
+
+
+def seed_all_backends(seed: Optional[int] = None) -> None:
+    """Sets all python, numpy and pytorch seeds."""
+
+    if seed is None:
+        seed = int(torch.randint(1_000_000, size=(1,)))
+
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True  # type: ignore
+    torch.backends.cudnn.benchmark = False  # type: ignore


Test patch:
diff --git a/tests/analysis_test.py b/tests/analysis_test.py
--- a/tests/analysis_test.py
+++ b/tests/analysis_test.py
@@ -64,7 +64,7 @@ def test_1d_marginals_peaks_from_kde(num_candidates):
 
     num_samples = 10_000
     num_dim = 4
-    samples = torch.ones(num_samples, num_dim) + 1.0 * torch.randn(num_samples, num_dim)
+    samples = torch.ones(num_samples, num_dim) + 0.5 * torch.randn(num_samples, num_dim)
 
     peaks = get_1d_marginal_peaks_from_kde(samples, num_candidates=num_candidates)
     assert torch.allclose(peaks, torch.ones(num_dim), atol=1e-1)
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,16 +1,16 @@
-import numpy
 import pytest
 import torch
 
-# Seed for `set_seed` fixture. Change to change random state of all seeded tests.
+from sbi.utils.sbiutils import seed_all_backends
+
+# Seed for `set_seed` fixture. Change to random state of all seeded tests.
 seed = 1
 
 
 # Use seed automatically for every test function.
 @pytest.fixture(autouse=True)
 def set_seed():
-    torch.manual_seed(seed)
-    numpy.random.seed(seed)
+    seed_all_backends(seed)
 
 
 @pytest.fixture(scope="session", autouse=True)
diff --git a/tests/linearGaussian_snle_test.py b/tests/linearGaussian_snle_test.py
--- a/tests/linearGaussian_snle_test.py
+++ b/tests/linearGaussian_snle_test.py
@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import pytest
-import torch
 from torch import eye, ones, zeros
 from torch.distributions import MultivariateNormal
 
@@ -83,7 +82,7 @@ def test_c2st_snl_on_linearGaussian():
 
     x_o = zeros(1, x_dim)
     num_samples = 1000
-    num_simulations = 3100
+    num_simulations = 3000
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(x_dim)
diff --git a/tests/simulator_utils_test.py b/tests/simulator_utils_test.py
--- a/tests/simulator_utils_test.py
+++ b/tests/simulator_utils_test.py
@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import pytest
+import torch
 from torch import ones, zeros
 
 from sbi.simulators.linear_gaussian import diagonal_linear_gaussian
@@ -14,18 +15,29 @@
 
 @pytest.mark.parametrize(
     "num_sims",
-    (0, 100, 1000),
+    (0, 1000),
 )
 @pytest.mark.parametrize("batch_size", (1, 100, 1000))
-@pytest.mark.parametrize("simulator", (diagonal_linear_gaussian, lambda _: ones(2)))
+@pytest.mark.parametrize(
+    "simulator", (diagonal_linear_gaussian, lambda _: torch.randn((2,)))
+)
+@pytest.mark.parametrize("seed", (None, 42))
 def test_simulate_in_batches(
     num_sims,
     batch_size,
     simulator,
+    seed,
     prior=BoxUniform(zeros(5), ones(5)),
 ):
     """Test combinations of num_sims and simulation_batch_size."""
 
     simulator, prior = prepare_for_sbi(simulator, prior)
     theta = prior.sample((num_sims,))
-    simulate_in_batches(simulator, theta, batch_size)
+    # run twice to check seeding.
+    x1 = simulate_in_batches(simulator, theta, batch_size, seed=seed)
+    x2 = simulate_in_batches(simulator, theta, batch_size, seed=seed)
+
+    if seed is None and num_sims > 0:
+        assert not torch.equal(x1, x2)
+    else:
+        assert torch.equal(x1, x2)