<jupyter_start><jupyter_text>Gaussian Mixture ModelThis is a brief tutorial on training mixture models in Pyro. We'll focus on the mechanics of `enum_discrete` and setting up mixture weights. To simplify matters, we'll train a trivial 1-D Gaussian model on a tiny 5-point dataset.<jupyter_code>import os
import sys
import torch
import numpy as np
import scipy.stats
from torch.nn import Softmax
from torch.nn.functional import softplus
from torch.autograd import Variable
from matplotlib import pyplot
%matplotlib inline

import pyro
from pyro.optim import Adam
from pyro.infer.svi import SVI
from pyro.distributions import categorical, normal
softmax = Softmax(dim=-1)
smoke_test = ('CI' in os.environ)<jupyter_output><empty_output><jupyter_text>DatasetHere is our tiny dataset. It has five points.<jupyter_code>data = Variable(torch.Tensor([0, 1, 10, 11, 12]))<jupyter_output><empty_output><jupyter_text>Maximum likelihood approachLet's start by optimizing model parameters `ps`, `mu`, and `sigma`, rather than treating them as random variables with priors. Our model will learn global mixture weights, the mean of each mixture component, and a shared variance that is common to both components. Our guide will learn soft assignment weights of each point. Note that none of our parameters have priors, so this approach is not fully Bayesian.<jupyter_code>K = 2  # Fixed number of components.

def model(data):
    # Global parameters.
    ps = softmax(pyro.param('unconstrained_ps', Variable(torch.zeros(K), requires_grad=True)))
    ps = ps.view(1, -1).expand(len(data), -1)
    mu = pyro.param('mu', Variable(torch.normal(torch.zeros(K), torch.ones(K)),
                                   requires_grad=True))
    sigma = softplus(pyro.param('unconstrained_sigma',
                                Variable(torch.Tensor([-1]), requires_grad=True)))

    # Local parameters.
    with pyro.iarange('data'):
        # Until Pyro 0.2, we'll have to .squeeze() and .unsqeeze() to satisfy iarange.
        z = pyro.sample('z', categorical, ps).squeeze(-1)
        mu_z = mu.index_select(0, z).unsqueeze(-1)
        pyro.sample('x', normal, mu_z, sigma.expand_as(mu_z), obs=data.unsqueeze(-1))

def guide(data):
    with pyro.iarange('data'):
        p = softmax(pyro.param('unconstrained_p',
                               Variable(torch.zeros(len(data), K), requires_grad=True)))
        pyro.sample('z', categorical, p)<jupyter_output><empty_output><jupyter_text>To run inference with this `(model,guide)` pair, we use Pyro's `enum_discrete` option to enumerate over all assignments in each iteration. Since we've wrapped the batched categorical assignments in a `pyro.iarange` indepencence context, this enumeration can happen in parallel: we enumerate only 2 possibilites, rather than `2**len(data) = 32`.<jupyter_code>optim = pyro.optim.Adam({'lr': 0.2, 'betas': [0.9, 0.99]})
inference = SVI(model, guide, optim, 'ELBO', enum_discrete=True)
pyro.set_rng_seed(0)            # Set seed to make results reproducible.
pyro.get_param_store().clear()  # Clear stale param values.
loss = []
for i in range(200 if not smoke_test else 2):
    inference.step(data)
    loss.append(inference.step(data))
    sys.stdout.write('.' if i % 80 else '\n')
    sys.stdout.flush()
pyplot.figure(figsize=(8,4)).set_facecolor('white')
pyplot.plot(loss)
pyplot.xlabel('iters')
pyplot.ylabel('loss')
pyplot.yscale('log')
pyplot.title('Convergence of SVI')<jupyter_output><empty_output><jupyter_text>Here are the learned parameters:<jupyter_code>ps = softmax(pyro.param('unconstrained_ps'))
mu = pyro.param('mu')
sigma = softplus(pyro.param('unconstrained_sigma'))
print('ps = {}'.format(ps.data.numpy()))
print('mu = {}'.format(mu.data.numpy()))
print('sigma = {}'.format(sigma.data.numpy()))<jupyter_output>ps = [ 0.5992229   0.40077707]
mu = [ 10.93902969   0.55254382]
sigma = [ 2.09332943]<jupyter_text>The model's `ps` are as expected, with 3/5 of the data in the first component and 2/3 in the second component. We can also examine the guide's local `p` variable which can be interpreted as assignment probabilities.<jupyter_code>p = softmax(pyro.param('unconstrained_p'))
pyplot.figure(figsize=(8, 4)).set_facecolor('white')
pyplot.plot(data.data.numpy(), p.data.numpy()[:,0], 'ro',
            label='component with mean {:0.2g}'.format(mu.data[0]))
pyplot.plot(data.data.numpy(), p.data.numpy()[:,1], 'bo',
            label='component with mean {:0.2g}'.format(mu.data[1]))
pyplot.title('Mixture assignment weights')
pyplot.xlabel('data value')
pyplot.ylabel('assignment weight')
pyplot.legend(loc='center');<jupyter_output><empty_output><jupyter_text>Next let's visualize the mixture model.<jupyter_code>X = np.arange(-5,20,0.1)
Y1 = ps.data[0] * scipy.stats.norm.pdf((X - mu.data[0]) / sigma.data[0])
Y2 = ps.data[1] * scipy.stats.norm.pdf((X - mu.data[1]) / sigma.data[0])

pyplot.figure(figsize=(8, 4)).set_facecolor('white')
pyplot.plot(X, Y1, 'r-')
pyplot.plot(X, Y2, 'b-')
pyplot.plot(X, Y1 + Y2, 'k--')
pyplot.plot(data.data.numpy(), np.zeros(len(data)), 'k*')
pyplot.title('Densitiy of two-component mixture model')
pyplot.ylabel('probability density');<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14203010420686993, 'fraction_numerical': 0.023157082207641837, 'mean_word_length': 4.714443219404631, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 28, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '295891', 'license': 'MIT', 'n_tokens_mistral': 1843, 'n_tokens_neox': 1722, 'n_words': 454, 'path': 'tutorial/source/gmm.ipynb', 'repo_name': 'DEVESHTARASIA/pyro', 'size': '5182'}