HELLO CYBERNETICS

深層学習、機械学習、強化学習、信号処理、制御工学、量子計算などをテーマに扱っていきます

【確率的プログラミング】Edward2, Pyro, PyStanのベイズ線形回帰コードメモ

 

 

follow us in feedly

f:id:s0sem0y:20181113233530p:plain

はじめに

特にベイズの説明もしませんし、コードの説明もしません。 単なるメモです。

共通のデータ

下記の関数より生成します。 内容は下記の記事と全く同じで、4つの会社から合計200人の社員の年収を22歳から45歳の年まで集めたデータという設定です。 今回の記事比較では、単に200人の22歳から45歳までの年収データという設定で線形回帰します。ちゃんと会社ごとの差を取り入れた階層ベイズモデリングを見たい方は下記の記事を覗いてみてください。

www.hellocybernetics.tech

N = 200
K = 4

def get_data():
    
    a0 = 350.
    b0 = 20.
    s_a = 40.
    s_b = 5.
    s_Y = 30.

    
    a = np.random.normal(loc=a0, scale=s_a, size=(K,))
    b = np.random.normal(loc=b0, scale=s_b, size=(K,))
    
    KID = []
    X = []
    Y = []
    
    for n in range(N):
        kid = np.random.randint(0, K, 1)
        KID.append(int(kid))
        
        x = np.random.randint(22, 45, 1)
        X.append(x)
        
        Y.append(np.random.normal(loc=a[kid]+b[kid]*(x-22), scale=s_Y))

    return (np.array(X).reshape(N, 1).astype(np.float32), 
            np.array(Y).reshape(N, 1), 
            np.array(KID).reshape(N, 1).astype(np.int32))

上が会社ごとに色を変えたデータで、下が一緒くたにしたデータ。今回は下のデータを使っていきます。

f:id:s0sem0y:20181113232205p:plain

f:id:s0sem0y:20181113232618p:plain

モデル

ベイズ線形回帰のモデルとしては下記の式を用います。

$$ \begin{align} p(y\mid x,a,b, \sigma) & = \mathcal N (a + b (x - 22), \sigma^2)\\ p(\sigma) &= \mathcal N(0, s) \\ p(a) &= \mathcal N(0, s) \\ p(b) &= \mathcal N(0, s) \end{align} $$

ここで$s$は十分大きな正の値としてパラメータ全てに無情報事前分布を敷いたという設定で行きます。

PyStan

モデル

データとして後で外から与える変数、パラメータ(事後分布を求めたい変数)、確率モデル、予測モデルをそれぞれ別々のブロックで書きます。かなり直感的で分かりやすいです。本当はベクトル化して書いたほうが良いでしょうけども、バッチサイズを気にせずに1つ1つのサンプリングを意識して書く書き方でも問題なく動きます。

stan_model = """
  data {
    int N;
    real X[N];
    real Y[N];
    int N_s;
    real X_s[N_s];
  }
  parameters {
    real a;
    real b;
    real<lower=0> sigma; 
  }
  model{
    for (n in 1:N){
      Y[n] ~ normal(a + b * (X[n] - 22), sigma);
    }
  }
  generated quantities {
    real Y_s[N_s];
    for (n in 1:N_s){
      Y_s[n] = normal_rng(a + b * (X_s[n] - 22), sigma);
    }
  }
"""

推論

X_data, Y_data, KID_data = get_data()
X_s = np.arange(22, 60, 1)
N_s = X_s.shape[0]

stan_data = {"N":X_data.shape[0], "X":X_data, "Y":Y_data, "N_s":N_s, "X_s":X_s}
sm = pystan.StanModel(model_code=stan_model)
fit = sm.sampling(data=stan_data, iter=2000, warmup=500, chains=3, seed=1992)

あとはpythonの世界でpystan.StanModelのインスタンスを作成します。このときに先程書いたstan_modelをわたします。 続いてpystan.StanModelsamplingメソッドにデータと取り出したいサンプル数などの情報を与えて動かすだけです。とても簡単。

事後分布

parametersブロックで書いた変数を下記のように取り出すことで、事後分布を近似したサンプルを取り出すことができます。numpy.arrayになっているので、あとは慣れた可視化ライブラリを使ってヒストグラムなりプロットなりを見ればいいでしょう。

ms_a = fit.extract("a")["a"]

Pyro

モデル

MCMCを用いて推論を行う場合は、def model内で使われる変数に関して、pyro.sample関数で得たものは「確率変数である」と認識され、その確率変数の事後分布を得ることができます。パラメータに限らず潜在変数でも何でも、とにかく事後分布を求めたいものは$pyro.sample$で表現してやればいいという簡単仕様です。

観測できる確率変数に関しては$pyro.sample(obs=observed_data)$として与えることで、パラメータや潜在変数と区別することができます。今回の場合、yは観測できる教師データです。

def model(x_data, y_data):
    loc, scale = 25.*torch.ones(1), 100.*torch.ones(1)
    weight = pyro.sample("weight",
        pyro.distributions.Normal(loc, scale)
    )
    
    bias_loc, bias_scale = 300.*torch.ones(1), 100. * torch.ones(1)
    bias = pyro.sample("bias",
        pyro.distributions.Normal(bias_loc, bias_scale)
    )
    
    scale = pyro.sample("scale",
        pyro.distributions.Normal(torch.zeros(1), 100. * torch.ones(1))
    )
    

    y = pyro.sample(
        "y",
        pyro.distributions.Normal(weight*(x_data-22) + bias, scale),
        obs=y_data
    )
    return y

推論

推論もべらぼうに簡単です。今回はNUTSを使います。 NUTSにモデルを渡してやることで、MCMCのカーネルとなるインスタンスを生成します。そして、MCMCのコンストラクタにカーネルと取り出したいサンプル数を渡して、runメソッドに観測データを投げてやるだけです(この時の投げる順番はmodelで指定した順番となっています)。

from pyro.infer.mcmc import MCMC, NUTS, HMC
from pyro.infer import EmpiricalMarginal, TracePredictive

x_data, y_data = torch.Tensor(X_data), torch.Tensor(Y_data)

nuts_kernel = NUTS(model, adapt_step_size=True)
mcmc_run = MCMC(
    nuts_kernel, 
    num_samples=2000, 
    warmup_steps=500).run(x_data, y_data)

事後分布

事後分布を取り出したい場合、EmpiricalMarginalクラスにmcmc_runを渡してやり、事後分布(を近似したサンプル)として取り出したい確率変数の名前をわたしてやるというのが一連の手順です。

posterior_w = EmpiricalMarginal(mcmc_run, 'weight')
posterior_b = EmpiricalMarginal(mcmc_run, 'bias')

取り出したposterior_wEmpiricalMarginalクラスのインスタンスとなっています。下記のメソッドでnumpy配列としてデータを取り出せるので、適当な可視化ライブラリでデータを見ることができます。

plt.figure()
sns.distplot(posterior_w.get_samples_and_weights()[0])
plt.figure()
sns.distplot(posterior_b.get_samples_and_weights()[0])

Edward2

モデル

Edward2もPyroにかなり似ています。pyro.sample関数のようなものはありませんが、edward2のAPIが提供するed.Normalなどは直接RandomVariavble型を吐き出すため、ed.***の分布によって吐き出された変数は自動的にMCMCの事後分布近似対象になります。

def linear_regression(features):
    weight = ed.Normal(loc=tf.zeros(features.shape[1]), scale=300., name="weight")
    bias = ed.Normal(loc=0., scale=300., name="bias")
    scale = ed.Normal(loc=0, scale=300., name="scale")
    
    linear_predictor = weight*(features-22) + bias
    outcomes = ed.Normal(
        loc=linear_predictor,
        scale=scale,
        name="outcomes")
    
    return outcomes

推論

edward2のAPIはココからが割としっかり手作業で記述が必要です。まず、(正規化されている必要はありませんが)対数同時確率を自分で準備しなければなりません。とはいっても、RandomVariable型をAPI側が認識できるため下記のコードでモデルをわたしてやるだけです。

log_joint = ed.make_log_joint_fn(linear_regression)

続いて、log_jointは高階関数となっているので、観測データであるfeaturesoutcomesを先に部分適用してしまいます。これでtarget_log_prob_fnはパラメータweight, bias, scaleのみを引数とする関数になっています(ココらへんはPyroでは隠蔽されている所なのでよく理解が必要です。HMCとその派生手法は事後分布を求めたいパラメータたちを引数とする関数の勾配を使った学習法です)。

features = tf.convert_to_tensor(X_data.reshape(-1, 1))
outcomes = tf.convert_to_tensor(Y_data.reshape(-1, 1))

def target_log_prob_fn(weight, bias, scale):
    return log_joint(features=features, 
                     weight=weight, 
                     bias=bias, 
                     scale=scale,
                     outcomes=outcomes)

さて、続いてまだ手作業で各部分はあります。とはいってもこれが最後で、APIにデータをわたしていくだけです。

まず、MCMCの初期値を決めます。ココを手で与えられるおかげなのか何なのか、わけのわからん初期値から始まったせいで対数確率が発散したりNaNになったりすることが避けられるのではないかなと思います(パラメータに対する知見があればですが)。

次に、遷移カーネルを選びます。TF-pにはNUTSは正式に入っていないので、HMCを使います。対数確率と、ステップサイズなどもろもろ渡しておきましょう。

最後にMCMCを走らせますが、tfp.mcmc.sample_chainの戻り値は

[ [prameter1, parameter2, parameter3], kernel_results ]となっているので、タプルで (parameters, kernel_results)と受け取っても良いですが、順番を忘れないうちに下記のように受け取るのがスタンダードなようです。しっかり初期値も与えてあげます。

weight_init = tf.constant([20.])
bias_init = tf.constant(300.)
scale_init = tf.constant(10.)

hmc = tfp.mcmc.HamiltonianMonteCarlo(
    target_log_prob_fn=target_log_prob_fn,
    num_leapfrog_steps=3,
    step_size=tf.Variable(0.01),
    step_size_update_fn=tfp.mcmc.make_simple_step_size_update_policy())

    
[
    [weight, bias, scale],
    kernel_results,
] = tfp.mcmc.sample_chain(
        num_results=2000,
        num_burnin_steps=500,
        current_state=[weight_init, bias_init, scale_init],
        parallel_iterations=500,
        kernel=hmc)

事後分布

上記のコードで受け取った

[weight, bias, scale]

こいつらがそれぞれ事後分布(を近似したサンプル)になっているので、好みの可視化ライブラリで見てみましょう。

それぞれのコード

pystan

practice_pystan/SalaryModeling.ipynb at master · hellocybernetics/practice_pystan · GitHub

pyro

Pyro_practice/Salarymodeling_pyro_nuts.ipynb at master · hellocybernetics/Pyro_practice · GitHub

edward2

Tensorflow-Probability-Tutorials/salary_modeling_edward2.ipynb at master · hellocybernetics/Tensorflow-Probability-Tutorials · GitHub