Batch sampling slow without warning (#1490)

* direct posterior: warning for slow batched sampling with large batch size

* add option 'importance' to sample_with argument

* bugfix: error message for batched sampling not implemented

* adapt threshold for warning

* change name to  for consistency with rejection.py

* change condition for acceptance rate warning in rejection sampling to total number of samples tried to draw instead of total number drawn

* update threshold for high batch size sampling

Author: dgedon

sbi/inference/posteriors/direct_posterior.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import warnings
 from typing import Optional, Union
 
 import torch
@@ -168,6 +169,17 @@ def sample_batched(
         num_samples = torch.Size(sample_shape).numel()
         condition_shape = self.posterior_estimator.condition_shape
         x = reshape_to_batch_event(x, event_shape=condition_shape)
+        num_xos = x.shape[0]
+
+        # throw warning if num_x * num_samples is too large
+        if num_xos * num_samples > 2**21:  # 2 million-ish
+            warnings.warn(
+                "Note that for batched sampling, the direct posterior sampling "
+                "generates {num_xos} * {num_samples} = {num_xos * num_samples} "
+                "samples. This can be slow and memory-intensive. Consider "
+                "reducing the number of samples or batch size.",
+                stacklevel=2,
+            )
 
         max_sampling_batch_size = (
             self.max_sampling_batch_size

sbi/inference/posteriors/importance_posterior.py

@@ -210,9 +210,9 @@ def sample_batched(
         show_progress_bars: bool = True,
     ) -> Tensor:
         raise NotImplementedError(
-            "Batched sampling is not implemented for ImportanceSamplingPosterior. \
-            Alternatively you can use `sample` in a loop \
-            [posterior.sample(theta, x_o) for x_o in x]."
+            "Batched sampling is not implemented for ImportanceSamplingPosterior. "
+            "Alternatively you can use `sample` in a loop "
+            "[posterior.sample(theta, x_o) for x_o in x]."
         )
 
     def _importance_sample(

sbi/inference/posteriors/vi_posterior.py

@@ -304,9 +304,9 @@ def sample_batched(
         show_progress_bars: bool = True,
     ) -> Tensor:
         raise NotImplementedError(
-            "Batched sampling is not implemented for VIPosterior. \
-            Alternatively you can use `sample` in a loop \
-            [posterior.sample(theta, x_o) for x_o in x]."
+            "Batched sampling is not implemented for VIPosterior. "
+            "Alternatively you can use `sample` in a loop "
+            "[posterior.sample(theta, x_o) for x_o in x]."
         )
 
     def log_prob(

sbi/inference/trainers/npe/npe_base.py

@@ -470,7 +470,7 @@ def build_posterior(
                 If `None`, use the latest neural density estimator that was trained.
             prior: Prior distribution.
             sample_with: Method to use for sampling from the posterior. Must be one of
-                [`direct` | `mcmc` | `rejection` | `vi`].
+                [`direct` | `mcmc` | `rejection` | `vi` | `importance`].
             mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
                 `hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
                 implementation of slice sampling; select `hmc`, `nuts` or `slice` for

sbi/samplers/rejection/rejection.py

@@ -321,7 +321,7 @@ def accept_reject_sample(
             max(int(1.5 * num_remaining / max(min_acceptance_rate, 1e-12)), 100),
         )
         if (
-            num_sampled_total.min().item() > 1000
+            num_samples_possible > 1000
             and min_acceptance_rate < warn_acceptance
             and not leakage_warning_raised
         ):