issue #568 : small embedding net as default

sbi/neural_nets/classifier.py

@@ -10,6 +10,8 @@
 from sbi.utils.sbiutils import standardizing_net, z_score_parser
 from sbi.utils.user_input_checks import check_embedding_net_device, check_data_device
 
+from sbi.utils.sbiutils import DefaultEmbeddingNet
+
 
 class StandardizeInputs(nn.Module):
     def __init__(self, embedding_net_x, embedding_net_y, dim_x, dim_y):
@@ -114,6 +116,11 @@ def build_linear_classifier(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net_x, DefaultEmbeddingNet):
+        embedding_net_x.build_network(batch_x)
+    if isinstance(embedding_net_y, DefaultEmbeddingNet):
+        embedding_net_y.build_network(batch_y)
 
     check_data_device(batch_x, batch_y)
     check_embedding_net_device(embedding_net=embedding_net_x, datum=batch_y)
@@ -165,6 +172,11 @@ def build_mlp_classifier(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net_x, DefaultEmbeddingNet):
+        embedding_net_x.build_network(batch_x)
+    if isinstance(embedding_net_y, DefaultEmbeddingNet):
+        embedding_net_y.build_network(batch_y)
 
     check_data_device(batch_x, batch_y)
     check_embedding_net_device(embedding_net=embedding_net_x, datum=batch_y)
@@ -224,6 +236,11 @@ def build_resnet_classifier(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net_x, DefaultEmbeddingNet):
+        embedding_net_x.build_network(batch_x)
+    if isinstance(embedding_net_y, DefaultEmbeddingNet):
+        embedding_net_y.build_network(batch_y)
 
     check_data_device(batch_x, batch_y)
     check_embedding_net_device(embedding_net=embedding_net_x, datum=batch_y)

sbi/neural_nets/flow.py

@@ -18,6 +18,8 @@
 from sbi.utils.torchutils import create_alternating_binary_mask
 from sbi.utils.user_input_checks import check_embedding_net_device, check_data_device
 
+from sbi.utils.sbiutils import DefaultEmbeddingNet
+
 
 def build_made(
     batch_x: Tensor = None,
@@ -51,6 +53,10 @@ def build_made(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net, DefaultEmbeddingNet):
+        embedding_net.build_network(batch_y)
+
     x_numel = batch_x[0].numel()
     # Infer the output dimensionality of the embedding_net by making a forward pass.
     check_data_device(batch_x, batch_y)
@@ -124,6 +130,10 @@ def build_maf(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net, DefaultEmbeddingNet):
+        embedding_net.build_network(batch_y)
+
     x_numel = batch_x[0].numel()
     # Infer the output dimensionality of the embedding_net by making a forward pass.
     check_data_device(batch_x, batch_y)
@@ -208,6 +218,10 @@ def build_nsf(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net, DefaultEmbeddingNet):
+        embedding_net.build_network(batch_y)
+
     x_numel = batch_x[0].numel()
     # Infer the output dimensionality of the embedding_net by making a forward pass.
     check_data_device(batch_x, batch_y)

sbi/neural_nets/mdn.py

@@ -10,6 +10,8 @@
 
 from sbi.utils.user_input_checks import check_embedding_net_device, check_data_device
 
+from sbi.utils.sbiutils import DefaultEmbeddingNet
+
 
 def build_mdn(
     batch_x: Tensor = None,
@@ -43,6 +45,10 @@ def build_mdn(
     Returns:
         Neural network.
     """
+    # Initialize default embedding net with linear layer
+    if isinstance(embedding_net, DefaultEmbeddingNet):
+        embedding_net.build_network(batch_y)
+
     x_numel = batch_x[0].numel()
     # Infer the output dimensionality of the embedding_net by making a forward pass.
     check_data_device(batch_x, batch_y)

sbi/utils/get_nn_models.py

@@ -14,14 +14,16 @@
 from sbi.neural_nets.flow import build_made, build_maf, build_nsf
 from sbi.neural_nets.mdn import build_mdn
 
+from sbi.utils.sbiutils import DefaultEmbeddingNet
+
 
 def classifier_nn(
     model: str,
     z_score_theta: Optional[str] = "independent",
     z_score_x: Optional[str] = "independent",
     hidden_features: int = 50,
-    embedding_net_theta: nn.Module = nn.Identity(),
-    embedding_net_x: nn.Module = nn.Identity(),
+    embedding_net_theta: nn.Module = DefaultEmbeddingNet(),
+    embedding_net_x: nn.Module = DefaultEmbeddingNet(),
 ) -> Callable:
     r"""
     Returns a function that builds a classifier for learning density ratios.
@@ -93,7 +95,7 @@ def likelihood_nn(
     hidden_features: int = 50,
     num_transforms: int = 5,
     num_bins: int = 10,
-    embedding_net: nn.Module = nn.Identity(),
+    embedding_net: nn.Module = DefaultEmbeddingNet(),
     num_components: int = 10,
 ) -> Callable:
     r"""
@@ -170,7 +172,7 @@ def posterior_nn(
     hidden_features: int = 50,
     num_transforms: int = 5,
     num_bins: int = 10,
-    embedding_net: nn.Module = nn.Identity(),
+    embedding_net: nn.Module = DefaultEmbeddingNet(),
     num_components: int = 10,
 ) -> Callable:
     r"""

sbi/utils/sbiutils.py

@@ -830,3 +830,31 @@ def gradient_ascent(
         return argmax_, max_val
 
     return theta_transform.inv(best_theta_overall), max_val
+
+
+class DefaultEmbeddingNet(nn.Module):
+    def __init__(self):
+        """Class for Default Embedding Net that maps the context onto the flow."""
+        super().__init__()
+        self.net = nn.Identity()
+
+    def build_network(self, input: Tensor):
+        """Builds Small Network with Linear Layer of same input/output dimensions
+            and a non-linear activation function (relu).
+        Args:
+            input: Data batch used to infer dimensionality (number of features)
+        """
+        num_features = input.shape[-1]
+        self.net = nn.Sequential(nn.Linear(num_features, num_features), nn.ReLU())
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Network forward pass.
+
+        Args:
+            x: Input tensor (batch_size, num_features)
+
+        Returns:
+            Network output (batch_size, num_features).
+        """
+        x = self.net(x)
+        return x

tests/inference_with_NaN_simulator_test.py

@@ -77,7 +77,7 @@ def test_inference_with_nan_simulator(
     likelihood_cov = 0.3 * eye(num_dim)
     x_o = zeros(1, num_dim)
     num_samples = 500
-    num_simulations = 2000
+    num_simulations = 5000
 
     def linear_gaussian_nan(
         theta, likelihood_shift=likelihood_shift, likelihood_cov=likelihood_cov

tests/linearGaussian_snle_test.py

@@ -83,7 +83,7 @@ def test_c2st_snl_on_linearGaussian(set_seed):
 
     x_o = zeros(1, x_dim)
     num_samples = 1000
-    num_simulations = 1000
+    num_simulations = 3100
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(x_dim)
@@ -145,7 +145,7 @@ def test_c2st_and_map_snl_on_linearGaussian_different(
         set_seed: fixture for manual seeding
     """
     num_samples = 500
-    num_simulations = 1000
+    num_simulations = 3000
     trials_to_test = [1]
 
     # likelihood_mean will be likelihood_shift+theta
@@ -239,7 +239,7 @@ def test_c2st_multi_round_snl_on_linearGaussian(num_trials: int, set_seed):
     num_dim = 2
     x_o = zeros((num_trials, num_dim))
     num_samples = 500
-    num_simulations_per_round = 500 * num_trials
+    num_simulations_per_round = 600 * num_trials
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(num_dim)

tests/linearGaussian_snpe_test.py

@@ -68,7 +68,7 @@ def test_c2st_snpe_on_linearGaussian(
 
     x_o = zeros(num_trials, num_dim)
     num_samples = 1000
-    num_simulations = 2500
+    num_simulations = 2600
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(num_dim)
@@ -292,7 +292,7 @@ def test_c2st_multi_round_snpe_on_linearGaussian(method_str: str, set_seed):
         ).set_default_x(x_o)
     elif method_str == "snpe_c":
         inference = SNPE_C(**creation_args)
-        theta, x = simulate_for_sbi(simulator, prior, 500, simulation_batch_size=50)
+        theta, x = simulate_for_sbi(simulator, prior, 900, simulation_batch_size=50)
         posterior_estimator = inference.append_simulations(theta, x).train()
         posterior1 = DirectPosterior(
             prior=prior, posterior_estimator=posterior_estimator
@@ -429,7 +429,7 @@ def simulator(theta):
     # We need a pretty big dataset to properly model the bimodality.
     theta, x = simulate_for_sbi(simulator, prior, 10000)
     posterior_estimator = inference.append_simulations(theta, x).train(
-        max_num_epochs=50
+        max_num_epochs=60
     )
 
     posterior = DirectPosterior(
@@ -513,7 +513,7 @@ def test_mdn_conditional_density(num_dim: int = 3, cond_dim: int = 1):
 
     x_o = zeros(1, num_dim)
     num_samples = 1000
-    num_simulations = 2500
+    num_simulations = 2600
     condition = 0.1 * ones(1, num_dim)
 
     dims = list(range(num_dim))

tests/linearGaussian_snre_test.py

@@ -80,7 +80,7 @@ def test_c2st_sre_on_linearGaussian(set_seed):
     x_dim = 2
     discard_dims = theta_dim - x_dim
     num_samples = 1000
-    num_simulations = 1000
+    num_simulations = 2100
 
     likelihood_shift = -1.0 * ones(
         x_dim
@@ -161,7 +161,7 @@ def test_c2st_sre_variants_on_linearGaussian(
 
     x_o = zeros(num_trials, num_dim)
     num_samples = 500
-    num_simulations = 2500 if num_trials == 1 else 40000
+    num_simulations = 2600 if num_trials == 1 else 40500
 
     # `likelihood_mean` will be `likelihood_shift + theta`.
     likelihood_shift = -1.0 * ones(num_dim)
@@ -278,7 +278,7 @@ def test_api_sre_sampling_methods(sampling_method: str, prior_str: str, set_seed
     num_dim = 2
     num_samples = 10
     num_trials = 2
-    num_simulations = 1000
+    num_simulations = 2100
     x_o = zeros((num_trials, num_dim))
     # Test for multiple chains is cheap when vectorized.
     num_chains = 3 if sampling_method == "slice_np_vectorized" else 1

tests/sbiutils_test.py

@@ -417,7 +417,7 @@ def test_z_scoring_structured():
     t = np.arange(0, 1, 0.001)
     x_sin = np.sin(t * 2 * torch.pi * 5)
     x = np.vstack([[(x_sin * (i + 1)) + (i * 2)] for i in range(10)])
-    t_batch = torch.tensor(x)
+    t_batch = torch.FloatTensor(x)
 
     #### API tests
     # Go through every permutation of options to test API.