Gaussian Mixture Models N.B.: covariance is assumed to be diagonal
| 17 | |
| 18 | |
| 19 | class GMM(nn.Module): |
| 20 | """ Gaussian Mixture Models |
| 21 | N.B.: covariance is assumed to be diagonal |
| 22 | """ |
| 23 | |
| 24 | def __init__(self, w, mu, sigma): |
| 25 | """ |
| 26 | p(x) = sum_i w[i] N(mu[i], sigma[i]^2 * I) |
| 27 | |
| 28 | config: |
| 29 | w: shape K X 1, mixture coefficients, must sum to 1 |
| 30 | mu: shape K X D, mean |
| 31 | sigma: shape K X D, (diagonal) variance |
| 32 | """ |
| 33 | super().__init__() |
| 34 | self.register_buffer('w', w) |
| 35 | self.register_buffer('mu', mu) |
| 36 | self.register_buffer('sigma', sigma) |
| 37 | self.K = w.shape[0] |
| 38 | self.D = mu.shape[1] |
| 39 | |
| 40 | @torch.no_grad() |
| 41 | def log_gaussian(self, x, mu, sigma): |
| 42 | """ log density of single (diagonal-covariance) multivariate Gaussian""" |
| 43 | return -0.5 * ((x - mu)**2 / sigma**2).sum(dim=1) - 0.5 * ( |
| 44 | self.D * np.log(2 * np.pi) + torch.log(torch.prod(sigma**2))) |
| 45 | |
| 46 | @torch.no_grad() |
| 47 | def log_prob(self, x): |
| 48 | return torch.logsumexp( |
| 49 | torch.stack([ |
| 50 | torch.log(self.w[kk]) + |
| 51 | self.log_gaussian(x, self.mu[kk], self.sigma[kk]) |
| 52 | for kk in range(self.K) |
| 53 | ]), 0) |
| 54 | |
| 55 | @torch.no_grad() |
| 56 | def sampling(self, num_samples): |
| 57 | m = torch.distributions.Categorical(self.w) |
| 58 | idx = m.sample((num_samples,)) |
| 59 | return self.mu[idx, :] + torch.randn(num_samples, self.D).to( |
| 60 | self.w.device) * self.sigma[idx, :] |
| 61 | |
| 62 | @torch.no_grad() |
| 63 | def langevin_sampling(self, x, num_steps=10, eta=1.0e+0, is_anneal=False): |
| 64 | eta_list = cosine_schedule(eta_max=eta, T=num_steps) |
| 65 | for ii in range(num_steps): |
| 66 | eta_ii = eta_list[ii] if is_anneal else eta |
| 67 | x = x.detach() |
| 68 | x.requires_grad = True |
| 69 | eng = -self.log_prob(x).sum() |
| 70 | grad = torch.autograd.grad(eng, x)[0] |
| 71 | x = x - eta_ii * grad + torch.randn_like(x) * np.sqrt(eta_ii * 2) |
| 72 | |
| 73 | return x.detach() |