Confused how to implement simple Bayesian NN

From the docs I gathered that to create a simple 2 hidden layer classifier with 3 inputs and 12 nodes in each layer, it looks like:

model = PyroModule[nn.Sequential](
    PyroModule[nn.Linear](3, 12),
    PyroModule[nn.Sigmoid](),
    PyroModule[nn.Linear](12, 12),
    PyroModule[nn.Sigmoid](),
    PyroModule[nn.Linear](12, 1),
    PyroModule[nn.Sigmoid]()
)
assert isinstance(model, nn.Sequential)
assert isinstance(model, PyroModule)

# Now we can be Bayesian about weights in the first layer.
model[0].weight = PyroSample(
    prior=dist.Normal(0,1).expand([3, 12]).to_event(2))
model[2].weight = PyroSample(
    prior=dist.Normal(0,1).expand([12, 12]).to_event(2))
model[4].weight = PyroSample(
    prior=dist.Normal(0,1).expand([12,1]).to_event(2))

I have no idea what the next step is after defining the network and can’t seem to find a single full example of such. I have made a few other Pyro models, but am new to the nn module.

Hi @thecity2, you can find a full example of how to train a BNN in Bayesian regression tutorial. Let me know if you need to clarify something.

1 Like

Thank you @fehiepsi . I have been going through that example, I am trying to use subsampling in the plate:

def forward(self, x, y=None):
    sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
    with pyro.plate("data", x.shape[0], subsample_size=10) as ind:
        mean = self.linear(x[ind]).squeeze(-1)
        obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y[ind])
    return mean

This gives the error:
RuntimeError Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/pyro/poutine/trace_messenger.py in call(self, *args, **kwargs)
164 try:
–> 165 ret = self.fn(*args, **kwargs)
166 except (ValueError, RuntimeError) as e:

56 frames
RuntimeError: t() expects a tensor with <= 2 dimensions, but self is 3D

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
RuntimeError: t() expects a tensor with <= 2 dimensions, but self is 3D
     Trace Shapes:         
      Param Sites:         
     Sample Sites:         
        sigma dist    |    
             value    |    
linear.weight dist 10 | 1 3
             value 10 | 1 3
  linear.bias dist 10 | 1  
             value 10 | 1  

What is the cause of this error?

I think performing self.linear(x[ind]) under plate data will add an additional data dimension to the weights. See this caution in Pyro Modules tutorial. One solution you can try is to add load_pyro_samples method and call it before the plate statement.

So what I ended up doing is using DataLoader to train the model, thus avoiding the issue. Would be nice to know how to do it with subsampling though.

using DataLoader to train the model

This is a better solution. Just make sure to scale your likelihood (see this tutorial for more explanation) with

with pyro.poutine.scale(scale=num_full_data / batch_size):
    obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y_batch)
1 Like

@fehiepsi I don’t see anything about pyro.poutine.scale in that link. Is that a replacement for using pyro.plate?

That tutorial will explain to you why we need to scale. You can use scale poutine as in my last comment.