【Python】Pyroでベイズ推定

最近確率的モデルを扱う[Pyro](https://pyro.ai/)を知り、面白そうだと思ったので、お試しとして触ってみました。
本記事は、そのソースコードの共有となります。ソースコードはJupyter Notebookで書いています。
理論的な説明はほぼありませんが、ご了承ください。

<details><summary>環境</summary>
Windows10
Python: 3.7.7
Jupyter Notebook: 1.0.0
PyTorch: 1.5.1
Pyro: 1.4.0
scipy: 1.5.2
numpy: 1.19.1
matplotlib: 3.3.0

</details>

また、githubにipynbファイルを公開しています。
https://github.com/isuya0508/exercise_Pyro/blob/master/bayesian_inference_bernulli.ipynb

## Pyroとは？

Pyroは、Pytorchをバックエンドにした確率的モデルを扱うライブラリです。
pipからインストールできます。

```
pip install pyro-ppl 
```

このとき、事前にPytorchのインストールが必要です。詳細は公式ページをご参照ください。

## 試しに動かしてみる

今回は、ベルヌーイ分布$Ber(p)$に従うデータから、そのパラメータ$p$を推定することを考えます。
まず、必要なモジュールをインポートします。

```Python
from collections import Counter

import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

import torch
import torch.distributions.constraints as constraints
import pyro
import pyro.distributions as dist
from pyro.optim import SGD, Adam
from pyro.infer import SVI, Trace_ELBO

%matplotlib inline

pyro.set_rng_seed(0)
pyro.enable_validation(True)
```

### データの生成

乱数でデータを作成します。このとき、型をPytorchのtensorにする必要があります。

```Python
obs = stats.bernoulli.rvs(0.7, size=30, random_state=1)
obs = torch.tensor(obs, dtype=torch.float32)
obs
```

```
> tensor([1., 0., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 0., 1., 1., 1., 1.,
          1., 1., 0., 0., 1., 1., 0., 0., 1., 1., 1., 0.])
```

データ中の1の個数を確認します。

```Python
Counter(obs.numpy())
```

```
> Counter({1.0: 23, 0.0: 7})
```

よって、パラメータ$p$の最尤推定量は$23/30 \fallingdotseq 0.77$です。[^2]
以降、$p$をベイズ推定していきます。

### 事前分布の設定

ベイズ推定では、パラメータの事前分布を仮定し、観測したデータを合わせて事後分布を求めます。

ベルヌーイ分布$Ber(p)$のパラメータ$p$については、事前分布をベータ分布を仮定するのが一般的です [^1]
Pyroでは、`model`メソッドで事前分布とデータのモデルを記述します。

```Python
def model(data):
    # 事前分布を仮定
    prior = dist.Beta(1, 1)

    # データのモデリング
    p = pyro.sample('p', prior)
    for i in range(len(data)):
        pyro.sample(f'obs{i}', dist.Bernoulli(p), obs=data[i])
```

今回は$Beta(1, 1)$を仮定していますが、実はこれは一様分布と一致します。

### 事後分布の設定

`guide`メソッドで事後分布を記述します。事後分布も、事前分布と同様にベータ分布とします。
このとき、事後分布のパラメータとして適当な初期値を与えます。

```Python
def guide(data):
    # 事後分布の定義
    alpha_q = pyro.param('alpha_q', torch.tensor(15), constraint=constraints.positive)
    beta_q = pyro.param('beta_q', torch.tensor(15), constraint=constraints.positive)
    posterior = dist.Beta(alpha_q, beta_q)

    pyro.sample('p', posterior)
```

この事後分布のパラメータ$\alpha, \beta$を求めることになります。

### 事後分布のフィッティング

事後分布のパラメータの推定方法について、今回は確率的変分推定を採用します。Pyroではこの方法を使うのが基本のようです。
今回のベルヌーイ分布の例では解析的に事後分布を求めることができるため、変分推定のような近似手法を用いるのはナンセンスですが、練習ということでこの方法を使ってみます。
（本来は、解析的に分布を求めることができない場合に使う手法ですね。）

実装について、勾配降下法(SGD)のインスタンスと、確率的変分推定(SVI)のインスタンスを生成します。
SVIのインスタンスを生成するときに、上記で定義した`model`と`guide`を与えます。

```Python
optimizer = SGD(dict(lr=0.0001, momentum=0.9))
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
```

あとは、`svi.step(obs)`を繰り返して事後分布のパラメータを更新していくだけです。

```Python
NUM_STEPS = 2000
pyro.clear_param_store()

history = {
    'loss': [],
    'alpha': [],
    'beta': []
}

for step in range(1, NUM_STEPS + 1):
    loss = svi.step(obs)
    
    history['loss'].append(loss)
    history['alpha'].append(pyro.param('alpha_q').item())
    history['beta'].append(pyro.param('beta_q').item())
    
    if step % 100 == 0:
        print(f'STEP: {step} LOSS: {loss}')

```

```
> 
STEP: 100 LOSS: 17.461310371756554
STEP: 200 LOSS: 18.102468490600586
(中略)
STEP: 1900 LOSS: 17.97793820500374
STEP: 2000 LOSS: 17.95139753818512
```

ここで、`history`にフィッティング中のLoss、事後分布のパラメータ$\alpha, \beta$を記録しています。
ステップごとのこれらの数値をプロットすると、次のようになります。（ソースコードは省略します。）

![ダウンロード .png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/558932/bd2e9763-bba9-9ce3-4730-9a6a8cdc4d80.png)!

最終的に得られた$\alpha, \beta$を確認し、事後分布の期待値、分散を確認します。

```Python
infered_alpha = pyro.param('alpha_q').item()
infered_beta = pyro.param('beta_q').item()
posterior = stats.beta(infered_alpha_beta, infered_beta_beta)

print(f'alpha: {infered_alpha}')
print(f'beta: {infered_beta}')
print(f'Expected: {posterior.expect()}')
print(f'Variance: {posterior.var()}')
```

```
>
alpha: 25.764650344848633
beta: 7.556574821472168
Expected: 0.7732203787899605
Variance: 0.005109101547631603
```

事前分布と事後分布をプロットしてみます。（ソースコードは省略します。）

![ダウンロード.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/558932/63517b79-10fd-9256-7d7e-996320d7ce83.png)

うまく推定できていそうです。

### まとめ
Pyroを使って簡単なベイズ推定を実行してみました。
今回は単純なベイズ推定でしたが、複雑なモデルを柔軟に、かつシンプルに記述できるのがPyroの強みかと思います。
公式ドキュメントには多くの確率的手法の例が載ってあり、これらを眺めているだけでも勉強になりそうです。


[^2]: $p$の最尤推定量は標本平均と一致します。
[^1]: $p$の事前分布をベータ分布にすることで、事後分布もベータ分布に従うことが知られています。このような分布を共役事前分布といいます。