diff --git a/docs/how_to_guide/20_time_series_embedding.ipynb b/docs/how_to_guide/20_time_series_embedding.ipynb index 9d219d84a..b8a820e66 100644 --- a/docs/how_to_guide/20_time_series_embedding.ipynb +++ b/docs/how_to_guide/20_time_series_embedding.ipynb @@ -80,9 +80,6 @@ "name": "stdout", "output_type": "stream", "text": [ - "(16, 10)\n", - "(64, 5)\n", - "(16, 2)\n", " Neural network successfully converged after 31 epochs." ] } diff --git a/sbi/neural_nets/embedding_nets/SC_embedding.py b/sbi/neural_nets/embedding_nets/SC_embedding.py index 8829ec7c3..837f26ae5 100644 --- a/sbi/neural_nets/embedding_nets/SC_embedding.py +++ b/sbi/neural_nets/embedding_nets/SC_embedding.py @@ -325,9 +325,8 @@ def forward(self, x: Tensor) -> Tensor: n_points = x.shape[1] - assert self.modes <= n_points // 2 + 1, ( - "Modes should be at most floor(n_points/2) + 1" - ) + if self.modes > n_points // 2 + 1: + raise ValueError("Modes should be at most floor(n_points/2) + 1") x = self.fc0(x) # (batch_size, n_points, in_channels) diff --git a/sbi/neural_nets/embedding_nets/causal_cnn.py b/sbi/neural_nets/embedding_nets/causal_cnn.py index e9827a0e1..18f38cc85 100644 --- a/sbi/neural_nets/embedding_nets/causal_cnn.py +++ b/sbi/neural_nets/embedding_nets/causal_cnn.py @@ -56,8 +56,7 @@ def WaveNetSRLikeAggregator( intermediate_channel_sizes: Optional[List] = None, stride_sizes: Union[int, List] = 1, ) -> nn.Module: - """ - Creates a non-causal 1D CNN aggregator based on the WaveNet speach recognition task + """Non-causal 1D CNN aggregator based on WaveNet's speech recognition setup. By default this function creates an aggregator with two CNN layers, after every convolution a maxpooling operation halves the number of timepoints. @@ -68,14 +67,14 @@ def WaveNetSRLikeAggregator( Args: in_channels: number of channels at input. num_timepoints: length of the input. - output dim: wanted number of features as output. + output_dim: wanted number of features as output. activation: activation to apply after the convolution. kernel_sizes: (optional) alter the kernel size used and the number of CNN layers (through the length of the kernel size vector). intermediate_channel_sizes: (optional) alter the intermediate channel sizes used, should have length = len(kernel_sizes) - 1. - stride_sizes = Optional alter the stride used, either a vector of - len = len(kernel_sizes) or a single integer, in which case the same stride + stride_sizes: (optional) alter the stride used, either a list of length + len(kernel_sizes) or a single integer, in which case the same stride is used in every convolution. Returns: @@ -93,22 +92,24 @@ def WaveNetSRLikeAggregator( ] if intermediate_channel_sizes is None: intermediate_channel_sizes = [64] - assert len(intermediate_channel_sizes) == len(kernel_sizes) - 1, ( - "Provided kernel size list should be exactly one element longer " - "than channel size list." - ) - intermediate_channel_sizes += [output_dim] - if isinstance(stride_sizes, List): - assert len(stride_sizes) == len(kernel_sizes), ( - "Provided stride size list should be have the same size as" - "the kernel size list." + if len(intermediate_channel_sizes) != len(kernel_sizes) - 1: + raise ValueError( + "intermediate_channel_sizes needs one entry less than kernel_sizes, but " + f"has {len(intermediate_channel_sizes)} entries for {len(kernel_sizes)} " + "kernel sizes." ) + intermediate_channel_sizes = [*intermediate_channel_sizes, output_dim] + if isinstance(stride_sizes, list): + if len(stride_sizes) != len(kernel_sizes): + raise ValueError( + f"stride_sizes has {len(stride_sizes)} entries, but kernel_sizes has " + f"{len(kernel_sizes)}. Pass one stride per kernel size." + ) else: stride_sizes = [stride_sizes] * len(kernel_sizes) non_causal_layers = [] for ll in range(len(kernel_sizes)): - print(aggregator_out_shape) conv_layer = nn.Conv1d( in_channels=in_channels if ll == 0 else intermediate_channel_sizes[ll - 1], out_channels=intermediate_channel_sizes[ll], @@ -131,7 +132,6 @@ def WaveNetSRLikeAggregator( / 2 ), ) - print(aggregator_out_shape) aggregator = nn.Sequential(*non_causal_layers, nn.AdaptiveAvgPool1d(1)) return aggregator @@ -152,10 +152,10 @@ def __init__( aggregator: Optional[nn.Module] = None, output_dim: int = 20, ): - """Intitialize embedding network that uses 1D causal convolutions. + """Initialize embedding network that uses 1D causal convolutions. This is a simplified version of the architecture introduced for - the speech recognition task in the WaveNet paper (van den Oord, et al. (2016)) + the speech recognition task in the WaveNet paper (van den Oord, et al. (2016)). After several dilated causal convolutions (that maintain the dimensionality of the input), an aggregator network is used to bring down the dimensionality. @@ -163,7 +163,7 @@ def __init__( If you do not provide an aggregator network yourself, a default aggregator is used. This default aggregator is based on the WaveNet paper's description of their Speech Recognition Task, and uses non-causal convolutions and pooling - layers, and global average poolingg to obtain a final low dimensional embedding. + layers, and global average pooling to obtain a final low dimensional embedding. Args: input_shape: Spatial input shape without batch or channel dimensions, @@ -191,17 +191,25 @@ def __init__( """ super(CausalCNNEmbedding, self).__init__() - assert isinstance(input_shape, Tuple), ( - "input_shape must be a Tuple of size 1, e.g. (timepoints,)." - ) - assert len(input_shape) == 1, "Currently only 1D causal CNNs are supported." + if not isinstance(input_shape, tuple): + raise TypeError( + "input_shape must be a tuple of length 1, e.g. (timepoints,), but " + f"got {type(input_shape).__name__}." + ) + if len(input_shape) != 1: + raise ValueError("Currently only 1D causal CNNs are supported.") self.input_shape = (in_channels, *input_shape) total_timepoints = input_shape[0] - assert total_timepoints >= pool_kernel_size, ( - "Please ensure that the pool kernel size is not " - "larger than the number of observed timepoints." - ) + if total_timepoints < pool_kernel_size: + raise ValueError( + f"pool_kernel_size ({pool_kernel_size}) must not be larger than the " + f"number of timepoints ({total_timepoints})." + ) + if num_conv_layers < 1: + raise ValueError( + f"num_conv_layers must be at least 1, but got {num_conv_layers}." + ) if isinstance(dilation, str): match dilation.lower(): case "exponential_cyclic": @@ -216,22 +224,37 @@ def __init__( dilation_per_layer = [1] * num_conv_layers case _: raise ValueError( - f"{dilation} is not a valid option, please use \"none\"," - "\"exponential\",or \"exponential_cyclic\", or pass a list " - "of dilation sizes." + f"dilation must be \"none\", \"exponential\", " + f"\"exponential_cyclic\" or a list of dilation sizes, but got " + f"\"{dilation}\"." ) else: - assert isinstance(dilation, List), ( - "Please pass dilation size as list or a string option." - ) + if not isinstance(dilation, list): + raise TypeError( + "dilation must be a list of dilation sizes or a string option, " + f"but got {type(dilation).__name__}." + ) + if len(dilation) != num_conv_layers: + raise ValueError( + f"dilation has {len(dilation)} entries, but num_conv_layers is " + f"{num_conv_layers}. Pass one dilation per convolutional layer." + ) dilation_per_layer = dilation - assert max(dilation_per_layer) < total_timepoints, ( - "Your maximal dilations size used is larger than the number of " - "timepoints in your input, please provide a list with smaller dilations." - ) + if max(dilation_per_layer) >= total_timepoints: + raise ValueError( + f"The largest dilation ({max(dilation_per_layer)}) must be smaller " + f"than the number of timepoints ({total_timepoints}). Pass smaller " + "dilations or fewer num_conv_layers." + ) if out_channels_per_layer is None: out_channels_per_layer = [16] * num_conv_layers + if len(out_channels_per_layer) != num_conv_layers: + raise ValueError( + f"out_channels_per_layer has {len(out_channels_per_layer)} entries, " + f"but num_conv_layers is {num_conv_layers}. Pass one entry per " + "convolutional layer." + ) causal_layers = [] for ll in range(num_conv_layers): @@ -255,10 +278,13 @@ def __init__( out_channels_per_layer[-1], int(total_timepoints / pool_kernel_size), ) - assert aggregator_out_shape[-1] > 1, ( - "Your dimensionality is already small," - "Please ensure a larger input size or use a custom aggregator." - ) + if aggregator_out_shape[-1] <= 1: + raise ValueError( + f"After pooling with pool_kernel_size={pool_kernel_size}, " + f"{aggregator_out_shape[-1]} timepoint(s) remain, but the default " + "aggregator needs at least 2. Use a smaller pool_kernel_size, a " + "longer input, or pass your own aggregator." + ) aggregator = WaveNetSRLikeAggregator( aggregator_out_shape[0], aggregator_out_shape[-1], diff --git a/sbi/neural_nets/embedding_nets/cnn.py b/sbi/neural_nets/embedding_nets/cnn.py index 7ce5e9eb6..ce7163469 100644 --- a/sbi/neural_nets/embedding_nets/cnn.py +++ b/sbi/neural_nets/embedding_nets/cnn.py @@ -105,46 +105,59 @@ def __init__( ): """Convolutional embedding network. - First two layers are convolutional, followed by fully connected layers. - - Automatically infers whether to apply 1D or 2D convolution depending on - input_shape. - Allows usage of multiple (color) channels by passing in_channels > 1. + Applies `num_conv_layers` blocks of convolution, ReLU and max pooling, + followed by fully connected layers. A 1D `input_shape` uses `Conv1d` and + `MaxPool1d`, a 2D `input_shape` uses `Conv2d` and `MaxPool2d`. Each + convolution (stride 1, padding 1) turns a spatial size `n` into + `n + 3 - kernel_size`, and each pooling step then divides it by + `pool_kernel_size`, rounding down. Args: input_shape: Spatial input shape without batch or channel dimensions, e.g., (28,) for 1D or (28, 28) for 2D. Inputs may be flat or use channel-first layout. - in_channels: Number of image channels, default 1. - out_channels_per_layer: Number of out convolutional out_channels for each - layer. Must match the number of layers passed below. - num_cnn_layers: Number of convolutional layers. - num_linear_layers: Number fully connected layer. - num_linear_units: Number of hidden units in fully-connected layers. - output_dim: Number of output units of the final layer. - kernel_size: Kernel size for both convolutional layers. - pool_size: pool size for MaxPool1d operation after the convolutional + in_channels: Number of input channels, e.g., 3 for RGB images. + out_channels_per_layer: Number of output channels of each convolutional + layer, one entry per layer. Defaults to 6 channels in the first + layer, doubling in each further layer, e.g., `[6, 12]` for two layers. + num_conv_layers: Number of convolutional layers. + num_linear_layers: Number of fully connected layers. + num_linear_units: Number of hidden units in the fully connected layers. + output_dim: Number of output units of the final layer. + kernel_size: Kernel size of every convolutional layer. + pool_kernel_size: Kernel size of the max pooling after every + convolutional layer. """ super(CNNEmbedding, self).__init__() - assert isinstance(input_shape, Tuple), ( - "input_shape must be a Tuple of size 1 or 2, e.g., (width, [height])." - ) - assert ( - 0 < len(input_shape) < 3 - ), """input_shape must be a Tuple of size 1 or 2, e.g., - (width, [height]). Number of input channels are passed separately""" + if not isinstance(input_shape, tuple): + raise TypeError( + "input_shape must be a tuple, e.g., (28,) or (28, 28), but got " + f"{type(input_shape).__name__}." + ) + if not 0 < len(input_shape) < 3: + raise ValueError( + "input_shape must have length 1 or 2, e.g., (28,) or (28, 28). " + "Pass the number of channels as in_channels." + ) use_2d_cnn = len(input_shape) == 2 conv_module = nn.Conv2d if use_2d_cnn else nn.Conv1d pool_module = nn.MaxPool2d if use_2d_cnn else nn.MaxPool1d + if num_conv_layers < 1: + raise ValueError( + f"num_conv_layers must be at least 1, but got {num_conv_layers}." + ) if out_channels_per_layer is None: - out_channels_per_layer = [6, 12] - assert len(out_channels_per_layer) == num_conv_layers, ( - "out_channels needs as many entries as num_cnn_layers." - ) + out_channels_per_layer = [6 * 2**ii for ii in range(num_conv_layers)] + if len(out_channels_per_layer) != num_conv_layers: + raise ValueError( + f"out_channels_per_layer has {len(out_channels_per_layer)} entries, " + f"but num_conv_layers is {num_conv_layers}. Pass one entry per " + "convolutional layer." + ) # define input shape with channel self.input_shape = (in_channels, *input_shape) @@ -156,7 +169,6 @@ def __init__( padding = 1 for ii in range(num_conv_layers): - # Defining another 2D convolution layer conv_layer = conv_module( in_channels=in_channels if ii == 0 else out_channels_per_layer[ii - 1], out_channels=out_channels_per_layer[ii], @@ -169,11 +181,15 @@ def __init__( # Calculate change of output size of each CNN layer cnn_output_size = get_new_cnn_output_size(cnn_output_size, conv_layer, pool) - assert all( - cnn_output_size - ), f"""CNN output size is zero at layer {ii + 1}. Either reduce - num_cnn_layers to {ii} or adjust the kernel_size - and pool_kernel_size accordingly.""" + if not all(size > 0 for size in cnn_output_size): + fewer_layers_hint = ( + f", or set num_conv_layers to at most {ii}" if ii > 0 else "" + ) + raise ValueError( + f"CNN output size {cnn_output_size} is not positive at layer " + f"{ii + 1}. Use a smaller kernel_size or pool_kernel_size" + f"{fewer_layers_hint}." + ) self.cnn_subnet = nn.Sequential(*cnn_layers) diff --git a/sbi/neural_nets/embedding_nets/lru.py b/sbi/neural_nets/embedding_nets/lru.py index 848a3acbb..714997709 100644 --- a/sbi/neural_nets/embedding_nets/lru.py +++ b/sbi/neural_nets/embedding_nets/lru.py @@ -337,10 +337,11 @@ def forward( ) else: state = state.to(device=input.device) - assert state.shape == expected_state_shape, ( - f"Invalid state shape {state.shape}, " # fmt: skip - f"expected {expected_state_shape}" # fmt: skip - ) + if state.shape != expected_state_shape: + raise ValueError( + f"Invalid state shape {state.shape}, " + f"expected {expected_state_shape}" + ) # Detemine which mode to run the forward method with. mode = mode or self.mode diff --git a/sbi/neural_nets/embedding_nets/permutation_invariant.py b/sbi/neural_nets/embedding_nets/permutation_invariant.py index e55a1e422..8db570d9a 100644 --- a/sbi/neural_nets/embedding_nets/permutation_invariant.py +++ b/sbi/neural_nets/embedding_nets/permutation_invariant.py @@ -70,10 +70,8 @@ def __init__( super().__init__() self.trial_net = trial_net self.aggregation_dim = aggregation_dim - assert aggregation_fn in [ - "mean", - "sum", - ], "aggregation_fn must be 'mean' or 'sum'." + if aggregation_fn not in ("mean", "sum"): + raise ValueError("aggregation_fn must be 'mean' or 'sum'.") self.aggregation_fn = aggregation_fn # construct fully connected layers @@ -121,6 +119,8 @@ def forward(self, x: Tensor) -> Tensor: x: Input tensor (batch_size, permutation_dim, input_dim) Returns: Network output (batch_size, output_dim). + Raises: + ValueError: If `x` contains inf, or the trial_net returns NaN or inf. """ if not getattr(self, "_warned_nan_padding", False): @@ -137,6 +137,12 @@ def forward(self, x: Tensor) -> Tensor: is_real_trial = ~torch.isnan(x).all(-1) trial_counts = is_real_trial.sum(dim=self.aggregation_dim).unsqueeze(-1) + # nan_to_num below would turn inf into a finite number. + if torch.isinf(x).any(): + raise ValueError( + "Input contains inf. Only NaN is allowed, to mark padded trials." + ) + # apply trial net with nan entries replaced with 0 masked_x = torch.nan_to_num(x, nan=0.0) trial_embeddings = self.trial_net(masked_x) @@ -144,15 +150,20 @@ def forward(self, x: Tensor) -> Tensor: trial_embeddings = trial_embeddings * is_real_trial.unsqueeze(-1).float() # Take mean over permutation dimension divide by number of trials - # (instead of just taking torch.mean) to account for masking. + # (instead of just taking torch.mean) to account for masking. Observations + # without real trials get a zero embedding, as with "sum". if self.aggregation_fn == "mean": - combined_embedding = ( - trial_embeddings.sum(dim=self.aggregation_dim) / trial_counts - ) + combined_embedding = trial_embeddings.sum( + dim=self.aggregation_dim + ) / trial_counts.clamp(min=1) else: combined_embedding = trial_embeddings.sum(dim=self.aggregation_dim) - assert not torch.isnan(combined_embedding).any(), "NaNs in embedding." + if not torch.isfinite(combined_embedding).all(): + raise ValueError( + "Non-finite value in the aggregated trial embedding: the trial_net " + "returned NaN or inf." + ) # add number of trials as additional input return self.fc_subnet(torch.cat([combined_embedding, trial_counts], dim=1)) diff --git a/tests/embedding_net_test.py b/tests/embedding_net_test.py index 7aaaea1b5..754f18b32 100644 --- a/tests/embedding_net_test.py +++ b/tests/embedding_net_test.py @@ -30,6 +30,7 @@ SpectralConvEmbedding, TransformerEmbedding, ) +from sbi.neural_nets.embedding_nets.causal_cnn import WaveNetSRLikeAggregator from sbi.neural_nets.embedding_nets.lru import LRU, LRUBlock from sbi.simulators.linear_gaussian import ( linear_gaussian, @@ -200,6 +201,64 @@ def test_cnn_embedding_input_layout(embedding_cls, input_shape, kwargs): embedding_net(channel_last_input) +@pytest.mark.parametrize( + ("input_shape", "kwargs", "match"), + [ + ((32, 32, 3), {}, "in_channels"), + ((32,), {"num_conv_layers": 0}, "num_conv_layers must be at least 1"), + ( + (32,), + {"num_conv_layers": 3, "out_channels_per_layer": [6, 12]}, + "num_conv_layers is 3", + ), + # Kernel much larger than the padded input: both sizes are negative, and + # their product is positive. + ((4, 4), {"kernel_size": 9}, "smaller kernel_size"), + ((8,), {"pool_kernel_size": 4}, "num_conv_layers to at most 1"), + ], +) +def test_cnn_embedding_errors_name_arguments(input_shape, kwargs, match): + """CNNEmbedding errors name the constructor arguments that fix them.""" + with pytest.raises(ValueError, match=match): + CNNEmbedding(input_shape, **kwargs) + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"num_conv_layers": 0}, "num_conv_layers must be at least 1"), + ({"out_channels_per_layer": [16, 16]}, "num_conv_layers is 5"), + ({"dilation": [1, 2]}, "dilation has 2 entries"), + ({"pool_kernel_size": 128}, r"pool_kernel_size \(128\) must not be larger"), + ({"pool_kernel_size": 64}, "the default aggregator needs at least 2"), + ], +) +def test_causal_cnn_embedding_errors_name_arguments(kwargs, match): + """CausalCNNEmbedding errors name the constructor arguments that fix them.""" + with pytest.raises(ValueError, match=match): + CausalCNNEmbedding((64,), **{"pool_kernel_size": 2, **kwargs}) + + +def test_wavenet_aggregator_keeps_arguments_and_prints_nothing(capsys): + """Building the aggregator neither changes the caller's list nor prints.""" + intermediate_channel_sizes = [8] + WaveNetSRLikeAggregator( + 4, 32, output_dim=5, intermediate_channel_sizes=intermediate_channel_sizes + ) + + assert intermediate_channel_sizes == [8] + assert capsys.readouterr().out == "" + + +def test_cnn_embedding_default_channels_follow_num_conv_layers(): + """The default channels double per layer, starting at 6.""" + embedding_net = CNNEmbedding((32,), num_conv_layers=3) + convs = [m for m in embedding_net.cnn_subnet if isinstance(m, nn.Conv1d)] + + assert [conv.out_channels for conv in convs] == [6, 12, 24] + assert embedding_net(torch.randn(4, 32)).shape == (4, 20) + + @pytest.mark.parametrize("input_shape", [(2,), (128,)]) @pytest.mark.parametrize("num_layers", [2, 4]) @pytest.mark.parametrize("num_hiddens", [32, 64]) @@ -571,15 +630,18 @@ def test_permutation_invariant_embedding_ignores_padding_per_observation( assert torch.allclose(batched[1], embedding(x[1:])[0]) +@pytest.mark.parametrize("aggregation_fn", ("mean", "sum")) @pytest.mark.parametrize( "pattern", ("padded_trial", "partially_nan_trial", "no_real_trial") ) -def test_permutation_invariant_embedding_warns_on_invalid_padding(pattern): +def test_permutation_invariant_embedding_warns_on_invalid_padding( + pattern, aggregation_fn +): """NaN outside whole padded trials must not silently become data; warn once.""" embedding = PermutationInvariantEmbedding( FCEmbedding(input_dim=2, output_dim=4), trial_net_output_dim=4, - aggregation_fn="sum", + aggregation_fn=aggregation_fn, ) x = torch.randn(2, 3, 2) x[0, 1:] = float("nan") @@ -590,11 +652,43 @@ def test_permutation_invariant_embedding_warns_on_invalid_padding(pattern): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") - embedding(x) + output = embedding(x) embedding(x) issued = [w for w in caught if "should only mark padded trials" in str(w.message)] assert len(issued) == (0 if pattern == "padded_trial" else 1) + if pattern == "no_real_trial": + # Zero aggregate and zero trial count, for "mean" as for "sum". + assert torch.allclose(output[1], embedding.fc_subnet(torch.zeros(1, 5))[0]) + + +@pytest.mark.parametrize("aggregation_fn", ("mean", "sum")) +@pytest.mark.parametrize("with_padding", (False, True)) +@pytest.mark.parametrize("inf", (float("inf"), float("-inf"))) +def test_permutation_invariant_embedding_rejects_inf(inf, with_padding, aggregation_fn): + """Inf in a real trial raises instead of passing as a huge finite embedding.""" + embedding = PermutationInvariantEmbedding( + FCEmbedding(input_dim=2, output_dim=4), + trial_net_output_dim=4, + aggregation_fn=aggregation_fn, + ) + x = torch.randn(2, 3, 2) + x[0, 1, 0] = inf + if with_padding: + x[:, 2] = float("nan") + + with pytest.raises(ValueError, match="Input contains inf"): + embedding(x) + + +def test_permutation_invariant_embedding_rejects_non_finite_trial_net_output(): + """A trial_net that returns inf for finite input raises.""" + embedding = PermutationInvariantEmbedding(nn.Linear(2, 4), trial_net_output_dim=4) + with torch.no_grad(): + embedding.trial_net.bias.fill_(float("inf")) + + with pytest.raises(ValueError, match="trial_net returned NaN or inf"): + embedding(torch.randn(2, 3, 2)) @pytest.mark.parametrize("input_shape", [(32, 32), (32, 64), (111, 111)])