Pyro 1.4 : Contributed : ガウス混合モデル (翻訳)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 08/08/2020 (1.4.0)

* 本ページは、Pyro の以下のドキュメントを翻訳した上で適宜、補足説明したものです :

• Contributed : Gaussian Mixture Model

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

ガウス混合モデル

これは混合モデルの動機づける例を通して Pyro で離散潜在変数をどのように周辺化するかを実演するチュートリアルです。tiny 5-ポイントデータセット上で自明な 1-D ガウスモデルを訓練することによりモデルを単純に維持して並列列挙の機構にフォーカスします。並列列挙へのより広範囲なイントロダクションについては enumeration チュートリアル も見てください。

import os
from collections import defaultdict
import torch
import numpy as np
import scipy.stats
from torch.distributions import constraints
from matplotlib import pyplot
%matplotlib inline

import pyro
import pyro.distributions as dist
from pyro import poutine
from pyro.infer.autoguide import AutoDelta
from pyro.optim import Adam
from pyro.infer import SVI, TraceEnum_ELBO, config_enumerate, infer_discrete

smoke_test = ('CI' in os.environ)
assert pyro.__version__.startswith('1.4.0')
pyro.enable_validation(True)

 

概要

Pyro の TraceEnum_ELBO はガイドとモデルの両者の変数を自動的に周辺化できます。ガイド変数を列挙するとき、Pyro はシーケンシャルに列挙するか (これは変数がダウンストリームに制御フローを決める場合に有用です)、変数のサンプルサイトでの可能な値の tensor を作成するために新しい tensor 次元を割当てて非標準的な評価を使用して並列に列挙することができます。それからこれらの非標準値はモデルで再生されます。モデルで変数を列挙するとき、変数は並列に列挙されなければならずそしてガイドで現れてはなりません。数学的には、モデル側列挙が変数を正確に周辺化することで Jensen の不等式の適用を回避する一方で、ガイド側列挙は総ての値を列挙することにより確率的 ELBO の分散を単純に減じます。

ここに tyny データセットがあります。それは 5 ポイントを持ちます。

data = torch.tensor([0., 1., 10., 11., 12.])

 

MAP 推定器を訓練する

事前分布とデータが与えられたときモデルパラメータ weights, locs と scale を学習することから始めましょう。AutoDelta ガイド (デルタ分布にちなんで命名されています) を使用してこれらの点推定を学習します。モデルはグローバル混合重み、各混合成分の位置、そして両者の成分に共通の共有スケールを学習します。推論の間、TraceEnum_ELBO はクラスタへのデータポイントの割当てを周辺化します。

K = 2  # Fixed number of components.

@config_enumerate
def model(data):
    # グローバル変数
    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))
    scale = pyro.sample('scale', dist.LogNormal(0., 2.))
    with pyro.plate('components', K):
        locs = pyro.sample('locs', dist.Normal(0., 10.))

    with pyro.plate('data', len(data)):
        # ローカル変数
        assignment = pyro.sample('assignment', dist.Categorical(weights))
        pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)

この (model, guide) ペアで推論を実行するため、各反復で総ての割当てに渡り列挙するため Pyro の config_enumerate() ハンドラを利用します。pyro.plate 独立性コンテキストでバッチ化 Categorical 割当てをラップしたので、この列挙は並列に発生することができます : 2**len(data) = 32 ではなく 2 つの可能性だけを列挙します。最後に、列挙の並列バージョンを使用するため、Pyro に max_plate_nesting=1 を通して単一 plate を使用しているだけであることを知らせます ; これは私達が右端次元 plate を使用していて Pyro は並列化のために任意の他の次元を利用できることを Pyro に知らせます。

optim = pyro.optim.Adam({'lr': 0.1, 'betas': [0.8, 0.99]})
elbo = TraceEnum_ELBO(max_plate_nesting=1)

推論の前に尤もらしい値に初期化します。混合モデルはローカルモードに非常に敏感です。一般的なアプローチは多くのランダムな初期化の中から最善を選択することです、そこではクラスタ平均はデータのランダム・サブサンプリングから初期化されます。AutoDelta ガイドを使用していますので、カスタム init_loc_fn() を定義することにより初期化できます。

def init_loc_fn(site):
    if site["name"] == "weights":
        # Initialize weights to uniform.
        return torch.ones(K) / K
    if site["name"] == "scale":
        return (data.var() / 2).sqrt()
    if site["name"] == "locs":
        return data[torch.multinomial(torch.ones(len(data)) / len(data), K)]
    raise ValueError(site["name"])

def initialize(seed):
    global global_guide, svi
    pyro.set_rng_seed(seed)
    pyro.clear_param_store()
    global_guide = AutoDelta(poutine.block(model, expose=['weights', 'locs', 'scale']),
                             init_loc_fn=init_loc_fn)
    svi = SVI(model, global_guide, optim, loss=elbo)
    return svi.loss(model, global_guide, data)

# Choose the best among 100 random initializations.
loss, seed = min((initialize(seed), seed) for seed in range(100))
initialize(seed)
print('seed = {}, initial_loss = {}'.format(seed, loss))

seed = 7, initial_loss = 25.665584564208984

訓練の間、収束を監視するため損失と勾配 norm の両者を集めます。これを PyTorch の .register_hook() メソッドを使用して行なうことができます。

# Register hooks to monitor gradient norms.
gradient_norms = defaultdict(list)
for name, value in pyro.get_param_store().named_parameters():
    value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))

losses = []
for i in range(200 if not smoke_test else 2):
    loss = svi.step(data)
    losses.append(loss)
    print('.' if i % 100 else '\n', end='')

pyplot.figure(figsize=(10,3), dpi=100).set_facecolor('white')
pyplot.plot(losses)
pyplot.xlabel('iters')
pyplot.ylabel('loss')
pyplot.yscale('log')
pyplot.title('Convergence of SVI');

pyplot.figure(figsize=(10,4), dpi=100).set_facecolor('white')
for name, grad_norms in gradient_norms.items():
    pyplot.plot(grad_norms, label=name)
pyplot.xlabel('iters')
pyplot.ylabel('gradient norm')
pyplot.yscale('log')
pyplot.legend(loc='best')
pyplot.title('Gradient norms during SVI');

ここに学習されたパラメータがあります :

map_estimates = global_guide(data)
weights = map_estimates['weights']
locs = map_estimates['locs']
scale = map_estimates['scale']
print('weights = {}'.format(weights.data.numpy()))
print('locs = {}'.format(locs.data.numpy()))
print('scale = {}'.format(scale.data.numpy()))

weights = [0.375 0.625]
locs = [ 0.49887404 10.984463  ]
scale = 0.6514337062835693

モデルの重みは予想どおり、最初の構成要素のデータのおよそ 2/5 そして 2 番目の構成要素の 3/5 です。次に混合モデルを可視化しましょう。

X = np.arange(-3,15,0.1)
Y1 = weights[0].item() * scipy.stats.norm.pdf((X - locs[0].item()) / scale.item())
Y2 = weights[1].item() * scipy.stats.norm.pdf((X - locs[1].item()) / scale.item())

pyplot.figure(figsize=(10, 4), dpi=100).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('Density of two-component mixture model')
pyplot.ylabel('probability density');

最後に混合モデルによる最適化は非凸でしばしば局所的な最適化条件で抜け出せなくなることに注意してください。例えばこのチュートリアルでは、scale が非常に大きく初期化される場合、everthing-in-one-cluster 仮説にスタックします。

 

モデルをサーブする: メンバーシップを予測する

混合モデルを訓練した今、モデルを分類器として利用することを望むかもしれません。訓練の間モデルの割当て変数を周辺化しました。これが高速な収束を提供する一方、それはガイドからのクラスタ割当てを読むことを妨げます。モデルを分類器として扱うために 2 つのオプションを議論します : 1 番目は infer_discrete を使用して (遥かに高速) そして 2 番目は SVI 内で列挙を使用して 2ndary ガイドを訓練することによります (遅いですがより一般的)。

 

離散的推論を使用してメンバーシップを予測する

メンバーシップを予測する最速の方法は trace と replay と一緒に infer_discrete ハンドラを使用することです。infer_discrete の temperature パラメータをゼロに設定して、MAP 分類器で始めましょう。trace, replay, and infer_discrete のような effect ハンドラのより深い見方は、effect ハンドラ・チュートリアル を見てください。

guide_trace = poutine.trace(global_guide).get_trace(data)  # record the globals
trained_model = poutine.replay(model, trace=guide_trace)  # replay the globals

def classifier(data, temperature=0):
    inferred_model = infer_discrete(trained_model, temperature=temperature,
                                    first_available_dim=-2)  # avoid conflict with data plate
    trace = poutine.trace(inferred_model).get_trace(data)
    return trace.nodes["assignment"]["value"]

print(classifier(data))

tensor([0, 0, 1, 1, 1])

実際に新しいデータ上でこの分類器を実行できます。

new_data = torch.arange(-3, 15, 0.1)
assignment = classifier(new_data)
pyplot.figure(figsize=(8, 2), dpi=100).set_facecolor('white')
pyplot.plot(new_data.numpy(), assignment.numpy())
pyplot.title('MAP assignment')
pyplot.xlabel('data value')
pyplot.ylabel('class assignment');

MAP 割当てでなくランダム事後分布割当てを生成するため、temperature=1 を設定できるでしょう。

print(classifier(data, temperature=1))

tensor([0, 0, 1, 1, 1])

クラスは非常に上手く分離されていますので、5.75 あたりで、クラス間の境界にズームインします。

new_data = torch.arange(5.5, 6.0, 0.005)
assignment = classifier(new_data, temperature=1)
pyplot.figure(figsize=(8, 2), dpi=100).set_facecolor('white')
pyplot.plot(new_data.numpy(), assignment.numpy(), 'bx', color='C0')
pyplot.title('Random posterior assignment')
pyplot.xlabel('data value')
pyplot.ylabel('class assignment');

 

ガイドで列挙することによりメンバーシップを予測する

クラス・メンバーシップを予測する 2 番目の方法はガイドで列挙することです。これは分類器モデルをサーブするためには上手く動作しません、何故ならば新しい各入力データバッチに対して確率的最適化を実行する必要があるからですが、それはより巨大な変分モデルに埋め込めるという点でより一般的です。

ガイドからクラスタ割当てを読むため、新しい full_guide を定義します、これは (上のような) グローバルパラメータと (前に周辺化された) ローカルパラメータの両者に fit します。グローバル変数のためには既に良い値学習していますので、poutine.block を使用して SVI をそれらを更新することからブロックします。

@config_enumerate
def full_guide(data):
    # Global variables.
    with poutine.block(hide_types=["param"]):  # Keep our learned values of global parameters.
        global_guide(data)

    # Local variables.
    with pyro.plate('data', len(data)):
        assignment_probs = pyro.param('assignment_probs', torch.ones(len(data), K) / K,
                                      constraint=constraints.unit_interval)
        pyro.sample('assignment', dist.Categorical(assignment_probs))

optim = pyro.optim.Adam({'lr': 0.2, 'betas': [0.8, 0.99]})
elbo = TraceEnum_ELBO(max_plate_nesting=1)
svi = SVI(model, full_guide, optim, loss=elbo)

# Register hooks to monitor gradient norms.
gradient_norms = defaultdict(list)
svi.loss(model, full_guide, data)  # Initializes param store.
for name, value in pyro.get_param_store().named_parameters():
    value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))

losses = []
for i in range(200 if not smoke_test else 2):
    loss = svi.step(data)
    losses.append(loss)
    print('.' if i % 100 else '\n', end='')

pyplot.figure(figsize=(10,3), dpi=100).set_facecolor('white')
pyplot.plot(losses)
pyplot.xlabel('iters')
pyplot.ylabel('loss')
pyplot.yscale('log')
pyplot.title('Convergence of SVI');

pyplot.figure(figsize=(10,4), dpi=100).set_facecolor('white')
for name, grad_norms in gradient_norms.items():
    pyplot.plot(grad_norms, label=name)
pyplot.xlabel('iters')
pyplot.ylabel('gradient norm')
pyplot.yscale('log')
pyplot.legend(loc='best')
pyplot.title('Gradient norms during SVI');

今ではガイドのローカル assignment_probs 変数を調べることができます。

assignment_probs = pyro.param('assignment_probs')
pyplot.figure(figsize=(8, 3), dpi=100).set_facecolor('white')
pyplot.plot(data.data.numpy(), assignment_probs.data.numpy()[:, 0], 'ro',
            label='component with mean {:0.2g}'.format(locs[0]))
pyplot.plot(data.data.numpy(), assignment_probs.data.numpy()[:, 1], 'bo',
            label='component with mean {:0.2g}'.format(locs[1]))
pyplot.title('Mixture assignment probabilities')
pyplot.xlabel('data value')
pyplot.ylabel('assignment probability')
pyplot.legend(loc='center');

 

MCMC

次に collapsed NUTS を使用して成分パラメータに渡る完全な事後分布を探求します、i.e. NUTS を使用して総ての離散的潜在変数を周辺化します。

from pyro.infer.mcmc.api import MCMC
from pyro.infer.mcmc import NUTS
pyro.set_rng_seed(2)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_samples=250, warmup_steps=50)
mcmc.run(data)
posterior_samples = mcmc.get_samples()

Sample: 100%|██████████| 300/300 [00:30<00:00,  9.99it/s, step size=1.69e-01, acc. rate=0.587]

X, Y = posterior_samples["locs"].t()

pyplot.figure(figsize=(8, 8), dpi=100).set_facecolor('white')
h, xs, ys, image = pyplot.hist2d(X.numpy(), Y.numpy(), bins=[20, 20])
pyplot.contour(np.log(h + 3).T, extent=[xs.min(), xs.max(), ys.min(), ys.max()],
               colors='white', alpha=0.8)
pyplot.title('Posterior density as estimated by collapsed NUTS')
pyplot.xlabel('loc of component 0')
pyplot.ylabel('loc of component 1')
pyplot.tight_layout()

混合成分の非同定可能性により、尤度 landscape は (11,0.5) と (0.5,11) の近くで 2 つの equally likely モードを持ちます。NUTS は 2 つのモードの間で切り替える難しさを持ちます。

pyplot.figure(figsize=(8, 3), dpi=100).set_facecolor('white')
pyplot.plot(X.numpy(), color='red')
pyplot.plot(Y.numpy(), color='blue')
pyplot.xlabel('NUTS step')
pyplot.ylabel('loc')
pyplot.title('Trace plot of loc parameter during NUTS inference')
pyplot.tight_layout()

以上