likelihood_nn embedding_net num_blocks

[tvwenger]

I am trying to use an embedding net with likelihood_nn to reduce my parameter dimensionality from 25 to 10. This fails with a RuntimeError: size mismatch.

Here is a minimum working example based on the example in the docs.

import torch
from torch import nn

from sbi.inference import SNLE, prepare_for_sbi
from sbi.utils.get_nn_models import likelihood_nn
import sbi.utils as utils

num_dim = 3
prior = utils.BoxUniform(low=-2*torch.ones(num_dim), 
                         high=2*torch.ones(num_dim))

def linear_gaussian(theta):
    return theta + 1.0 + torch.randn_like(theta) * 0.1

simulator, prior = prepare_for_sbi(linear_gaussian, prior)

embedding_net = nn.Sequential(
    nn.Linear(num_dim, 3),
    nn.ReLU(),
    nn.Linear(3, 3),
    nn.ReLU(),
    nn.Linear(3, 2))

density_estimator_build_fun = likelihood_nn(
    model='nsf', hidden_features=60, num_transforms=3, embedding_net=embedding_net)

inference = SNLE(simulator, prior, density_estimator=density_estimator_build_fun, 
                 show_progress_bars=False)

posterior = inference(num_simulations=500)

And the traceback:

Traceback (most recent call last):
  File "sbi_test.py", line 30, in <module>
    posterior = inference(num_simulations=500)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/inference/snle/snle_a.py", line 128, in __call__
    return super().__call__(**kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/inference/snle/snle_base.py", line 171, in __call__
    discard_prior_samples=discard_prior_samples,
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/inference/snle/snle_base.py", line 259, in _train
    log_prob = self._posterior.net.log_prob(theta_batch, context=x_batch)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/distributions/base.py", line 40, in log_prob
    return self._log_prob(inputs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/flows/base.py", line 39, in _log_prob
    noise, logabsdet = self._transform(inputs, context=embedded_context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 56, in forward
    return self._cascade(inputs, funcs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 56, in forward
    return self._cascade(inputs, funcs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 56, in forward
    return self._cascade(inputs, funcs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/base.py", line 50, in _cascade
    outputs, logabsdet = func(outputs, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/coupling.py", line 83, in forward
    transform_params = self.transform_net(identity_split, context)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/nn/nets/resnet.py", line 96, in forward
    temps = self.initial_layer(torch.cat((inputs, context), dim=1))
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 91, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py", line 1674, in linear
    ret = torch.addmm(bias, input, weight.t())
RuntimeError: size mismatch, m1: [50 x 3], m2: [4 x 60] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41

[michaeldeistler]

You are probably using sbi version 0.11.x. This was a bug that got fixed in 0.12.0. Does updating sbi fix it?

[tvwenger]

You are right! I updated to 0.12.0, but now I get a whole new error!

Traceback (most recent call last):
  File "sbi_test.py", line 30, in <module>
    posterior = inference(num_simulations=500)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/inference/snle/snle_a.py", line 128, in __call__
    return super().__call__(**kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/inference/snle/snle_base.py", line 154, in __call__
    neural_net=self._build_neural_net(theta, x),
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/utils/get_nn_models.py", line 153, in build_fn
    return build_nsf(batch_x=batch_x, batch_y=batch_theta, **kwargs)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/neural_nets/flow.py", line 204, in build_nsf
    for i in range(num_transforms)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/neural_nets/flow.py", line 204, in <listcomp>
    for i in range(num_transforms)
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/coupling.py", line 460, in __init__
    unconditional_transform=unconditional_transform,
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/nflows/transforms/coupling.py", line 53, in __init__
    self.num_transform_features * self._transform_dim_multiplier(),
  File "/home/twenger/anaconda3/lib/python3.7/site-packages/sbi/neural_nets/flow.py", line 191, in <lambda>
    num_blocks=num_blocks,
NameError: name 'num_blocks' is not defined

[michaeldeistler]

Ok, damn, this is a major bug from our side. I'll fix it asap and make a new release v0.12.1. I'll let you know once it's on PyPI.

If you can't wait, the fix is already in the main branch. So git clone https://github.com/mackelab/sbi.git and pip install . will give you access to a working version.

Thanks for reporting this!!!

[tvwenger]

Got it, thanks!

[janfb]

The new release with the fix is out. Thanks again for reporting!