fix: Update RatioEstimator classifier argument to use a Protocol (#1582)

* fix: Update RatioEstimator classifier argument to use a Protocol

* test: add tests for RatioEstimatorBuilder protocol

* Combine NRE classifier builder tests into a single function

* Fix formatting and import sorting

* Remove alias import for RatioEstimator

Author: abelaba

sbi/inference/trainers/nre/bnre.py

@@ -1,13 +1,14 @@
 # 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/>
 
-from typing import Callable, Dict, Optional, Union
+from typing import Dict, Optional, Union
 
 import torch
 from torch import Tensor, nn, ones
 from torch.distributions import Distribution
 
 from sbi.inference.trainers.nre.nre_a import NRE_A
+from sbi.inference.trainers.nre.nre_base import RatioEstimatorBuilder
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.torchutils import assert_all_finite
@@ -28,7 +29,7 @@ class BNRE(NRE_A):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, Callable] = "resnet",
+        classifier: Union[str, RatioEstimatorBuilder] = "resnet",
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[TensorboardSummaryWriter] = None,
@@ -42,11 +43,11 @@ def __init__(
                 prior must be passed to `.build_posterior()`.
             classifier: Classifier trained to approximate likelihood ratios. If it is
                 a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet). Alternatively, a function that builds a custom
-                neural network can be provided. The function will be called with the
-                first batch of simulations $(\theta, x)$, which can thus be used for
-                shape inference and potentially for z-scoring. It needs to return a
-                PyTorch `nn.Module` implementing the classifier.
+                linear, mlp, resnet), or a callable that implements the
+                `RatioEstimatorBuilder` protocol. The callable will be called with the
+                first batch of simulations (theta, x), which can thus be used for
+                shape inference and potentially for z-scoring. It returns a
+                `RatioEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_a.py

@@ -1,13 +1,16 @@
 # 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/>
 
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import torch
 from torch import Tensor, nn, ones
 from torch.distributions import Distribution
 
-from sbi.inference.trainers.nre.nre_base import RatioEstimatorTrainer
+from sbi.inference.trainers.nre.nre_base import (
+    RatioEstimatorBuilder,
+    RatioEstimatorTrainer,
+)
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.torchutils import assert_all_finite
@@ -23,7 +26,7 @@ class NRE_A(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, Callable] = "resnet",
+        classifier: Union[str, RatioEstimatorBuilder] = "resnet",
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[TensorboardSummaryWriter] = None,
@@ -37,11 +40,11 @@ def __init__(
                 prior must be passed to `.build_posterior()`.
             classifier: Classifier trained to approximate likelihood ratios. If it is
                 a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet). Alternatively, a function that builds a custom
-                neural network can be provided. The function will be called with the
-                first batch of simulations (theta, x), which can thus be used for shape
-                inference and potentially for z-scoring. It needs to return a PyTorch
-                `nn.Module` implementing the classifier.
+                linear, mlp, resnet), or a callable that implements the
+                `RatioEstimatorBuilder` protocol. The callable will be called with the
+                first batch of simulations (theta, x), which can thus be used for
+                shape inference and potentially for z-scoring. It returns a
+                `RatioEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_b.py

@@ -1,13 +1,16 @@
 # 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/>
 
-from typing import Callable, Dict, Optional, Union
+from typing import Dict, Optional, Union
 
 import torch
 from torch import Tensor, nn
 from torch.distributions import Distribution
 
-from sbi.inference.trainers.nre.nre_base import RatioEstimatorTrainer
+from sbi.inference.trainers.nre.nre_base import (
+    RatioEstimatorBuilder,
+    RatioEstimatorTrainer,
+)
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.torchutils import assert_all_finite
@@ -23,7 +26,7 @@ class NRE_B(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, Callable] = "resnet",
+        classifier: Union[str, RatioEstimatorBuilder] = "resnet",
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[TensorboardSummaryWriter] = None,
@@ -37,11 +40,11 @@ def __init__(
                 prior must be passed to `.build_posterior()`.
             classifier: Classifier trained to approximate likelihood ratios. If it is
                 a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet). Alternatively, a function that builds a custom
-                neural network can be provided. The function will be called with the
-                first batch of simulations (theta, x), which can thus be used for shape
-                inference and potentially for z-scoring. It needs to return a PyTorch
-                `nn.Module` implementing the classifier.
+                linear, mlp, resnet), or a callable that implements the
+                `RatioEstimatorBuilder` protocol. The callable will be called with the
+                first batch of simulations (theta, x), which can thus be used for
+                shape inference and potentially for z-scoring. It returns a
+                `RatioEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_base.py

@@ -4,7 +4,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Optional, Protocol, Union
 
 import torch
 from torch import Tensor, eye, nn, ones
@@ -18,6 +18,7 @@
 from sbi.inference.potentials import ratio_estimator_based_potential
 from sbi.inference.trainers.base import NeuralInference
 from sbi.neural_nets import classifier_nn
+from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.utils import (
     check_estimator_arg,
     check_prior,
@@ -26,11 +27,28 @@
 from sbi.utils.torchutils import repeat_rows
 
 
+class RatioEstimatorBuilder(Protocol):
+    """Protocol for building a ratio estimator from data."""
+
+    def __call__(self, theta: Tensor, x: Tensor) -> RatioEstimator:
+        """Build a ratio estimator from theta and x, which mainly inform the
+        shape of the input and the condition to the neural network.
+
+        Args:
+            theta: Parameter sets.
+            x: Simulation outputs.
+
+        Returns:
+            Ratio Estimator.
+        """
+        ...
+
+
 class RatioEstimatorTrainer(NeuralInference, ABC):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, Callable] = "resnet",
+        classifier: Union[str, RatioEstimatorBuilder] = "resnet",
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -56,11 +74,11 @@ def __init__(
         Args:
             classifier: Classifier trained to approximate likelihood ratios. If it is
                 a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet). Alternatively, a function that builds a custom
-                neural network can be provided. The function will be called with the
-                first batch of simulations (theta, x), which can thus be used for shape
-                inference and potentially for z-scoring. It needs to return a PyTorch
-                `nn.Module` implementing the classifier.
+                linear, mlp, resnet), or a callable that implements the
+                `RatioEstimatorBuilder` protocol. The callable will be called with the
+                first batch of simulations (theta, x), which can thus be used for
+                shape inference and potentially for z-scoring. It returns a
+                `RatioEstimator`.
 
         See docstring of `NeuralInference` class for all other arguments.
         """

sbi/inference/trainers/nre/nre_c.py

@@ -1,13 +1,16 @@
 # 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/>
 
-from typing import Callable, Dict, Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from torch import Tensor, nn
 from torch.distributions import Distribution
 
-from sbi.inference.trainers.nre.nre_base import RatioEstimatorTrainer
+from sbi.inference.trainers.nre.nre_base import (
+    RatioEstimatorBuilder,
+    RatioEstimatorTrainer,
+)
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.torchutils import assert_all_finite
@@ -37,7 +40,7 @@ class NRE_C(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, Callable] = "resnet",
+        classifier: Union[str, RatioEstimatorBuilder] = "resnet",
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[TensorboardSummaryWriter] = None,
@@ -51,11 +54,11 @@ def __init__(
                 prior must be passed to `.build_posterior()`.
             classifier: Classifier trained to approximate likelihood ratios. If it is
                 a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet). Alternatively, a function that builds a custom
-                neural network can be provided. The function will be called with the
-                first batch of simulations (theta, x), which can thus be used for shape
-                inference and potentially for z-scoring. It needs to return a PyTorch
-                `nn.Module` implementing the classifier.
+                linear, mlp, resnet), or a callable that implements the
+                `RatioEstimatorBuilder` protocol. The callable will be called with the
+                first batch of simulations (theta, x), which can thus be used for
+                shape inference and potentially for z-scoring. It returns a
+                `RatioEstimator`.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

tests/ratio_estimator_test.py

@@ -5,12 +5,14 @@
 
 import pytest
 import torch
-from torch import eye, zeros
+from torch import Tensor, eye, zeros
 from torch.distributions import MultivariateNormal
 
+from sbi.inference import NRE
 from sbi.neural_nets.embedding_nets import CNNEmbedding
 from sbi.neural_nets.net_builders import build_linear_classifier
 from sbi.neural_nets.ratio_estimators import RatioEstimator
+from sbi.utils.torchutils import BoxUniform
 
 
 class EmbeddingNet(torch.nn.Module):
@@ -72,3 +74,57 @@ def test_api_ratio_estimator(ratio_estimator, theta_shape, x_shape):
         nsamples,
     ), f"""unnormalized_log_ratio shape is not correct. It is of shape
     {unnormalized_log_ratio.shape}, but should be {(nsamples,)}"""
+
+
+def build_classifier(theta, x):
+    net = torch.nn.Linear(theta.shape[1] + x.shape[1], 1)
+    return RatioEstimator(net=net, theta_shape=theta[0].shape, x_shape=x[0].shape)
+
+
+def build_classifier_missing_args():
+    pass
+
+
+def build_classifier_missing_return(theta: Tensor, x: Tensor):
+    pass
+
+
+@pytest.mark.parametrize(
+    "classifier_builder",
+    [
+        build_classifier,
+        pytest.param(
+            build_classifier_missing_args,
+            marks=pytest.mark.xfail(
+                raises=TypeError,
+                reason="Missing required parameters in classifier builder.",
+            ),
+        ),
+        pytest.param(
+            build_classifier_missing_return,
+            marks=pytest.mark.xfail(
+                raises=AttributeError,
+                reason="Missing return of RatioEstimator in classifier builder.",
+            ),
+        ),
+    ],
+)
+def test_nre_with_valid_and_invalid_classifier_builders(classifier_builder):
+    r"""Test NRE works with valid classifier builders and fails with invalid ones.
+
+    Args:
+        classifier_builder: Function to build the classifier.
+    """
+
+    def simulator(theta):
+        return 1.0 + theta + torch.randn(theta.shape, device=theta.device) * 0.1
+
+    num_dim = 3
+    prior = BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))
+    theta = prior.sample((300,))
+    x = simulator(theta)
+
+    inference = NRE(classifier=classifier_builder)
+    inference.append_simulations(theta, x)
+
+    inference.train(max_num_epochs=1)