I am aware of pyro facilitating probabilistic models through standard SVI inference. But is it possible to write Bayesian models in pure pytorch? Say for instance, MAP training in Bayesian GMM.

I specify a bunch of priors and a likelihood, provide a MAP objective and learn point estimates but I am missing something key in my attempt here, perhaps whats confusing is that means, scales and weights have priors but they still need to be trainable as params in MAP.

```
import numpy as np
import torch
from tqdm import trange
import torch.distributions as D
from matplotlib.patches import Ellipse
import matplotlib.pyplot as plt
import sklearn.datasets as skd
class BayesianGMM(torch.nn.Module):
def __init__(self, Y, K, dim=2):
'''
:param Y (torch.tensor): data
:param K (int): number of fixed components
:param D (int): the dimension of data space.
'''
super(BayesianGMM, self).__init__()
self.Y = Y
self.K = K
self.dim = dim
self.mu_prior = torch.zeros((self.K,self.dim))
self.sig_prior = torch.ones((self.K, self.dim))*3.0
self.concentration = torch.randint(high=100,size=(1,))
self.alpha = torch.nn.Parameter(torch.tensor([self.concentration]*self.K).float())
self.means = D.Normal(self.mu_prior, self.sig_prior).sample()
self.scales = D.Gamma(torch.tensor([1.0]), torch.tensor([1.0])).sample()
self.weights = D.Dirichlet(self.alpha).sample()
self.z = D.Categorical(self.weights)
def forward(self, Y):
comp = D.Independent(D.Normal(self.means, self.scales), 1)
return D.MixtureSameFamily(self.z, comp).log_prob(Y)
def log_prior(self, Y):
self.means.log_prob(Y) + self.scales.log_prob(Y) + self.weights.log_prob(Y)
def get_trainable_param_names(self):
''' Prints a list of parameters which will be
learnt in the process of optimising the objective '''
for name, value in self.named_parameters():
print(name)
def loglikelihood(self, Y):
return self.forward(Y).mean()
def map_objective(self, Y):
return self.loglikelihood(Y) + self.log_prior(Y)
def train(self, Y, optimizer, n_steps):
losses = np.zeros(n_steps)
bar = trange(n_steps, leave=True)
for step in bar:
optimizer.zero_grad()
loss = self.map_objective(Y)
loss.backward()
optimizer.step()
losses[step] = loss
return losses
if __name__ == '__main__':
# Generate some data
Y, labels = skd.make_blobs(n_samples=2000, random_state=42,
cluster_std=[2.5, 0.5, 1.0])
Y = torch.tensor(Y).float()
model = BayesianGMM(Y, K=3)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
losses = model.train(Y, optimizer, 4000)
```