How to get prediction scores with Pyro models?

10

My reply on the other thread (How to avoid this error?) is helpful here too (unsurprisingly given I was trying to achieve the same thing). All I’ve added here is the Predictive bit again from the tutorials.

neural_network = nn.Sequential(
     nn.Linear(28 * 28, 100),
     nn.Sigmoid(),
     nn.Linear(100, 100),
     nn.Sigmoid(),
     nn.Linear(100, 1),
 )
 
module.to_pyro_module_(neural_network)
 
for m in neural_network.modules():
     for name, value in list(m.named_parameters(recurse=False)):
         setattr(m, name, module.PyroSample(prior=dist.Normal(0, 1)
                                       .expand(value.shape)
                                       .to_event(value.dim())))

This bit is important

class BayesianNeuralNetwork(PyroModule):
     def __init__(self, neural_network):
         super().__init__()
        self.neural_network= neural_network
 
     def forward(self, x, y=None):
         sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
         mean = self.neural_network(x).squeeze(-1)
         with pyro.plate("data", x.shape[0]):
             obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
         return mean

model = BayesianNeuralNetwork(neural_network)
 
guide= guides.AutoDiagonalNormal(model)
 
optimizer = Adam({"lr": 0.03}) 
 
svi= SVI(model, guide, optimizer, loss=Trace_ELBO())
 
X=torch.rand(85,10)
y=torch.rand(85,1)

pyro.clear_param_store()
svi.step(X,y)

predictive = Predictive(model, guide=guide, num_samples=100,return_sites=("obs", "_RETURN"))
samples = predictive(data['data'].float().to(device))
mean=samples['_RETURN'].mean(dim=0)
sigma = samples['_RETURN'].std(dim=0)

Where mean is your prediction and sigma is your standard deviation of the predictions. Hope this helps @h56cho