GPyTorch 1.2 Examples : Exact GP (回帰) – Fully Bayesian GP : NUTS でハイパーパラメータをサンプリング

by

GPyTorch 1.2 Examples : Exact GP (回帰) – Fully Bayesian GP : NUTS でハイパーパラメータをサンプリング (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 09/05/2020 (1.2)

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

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

 

Exact GP (回帰) – Fully Bayesian GP : NUTS でハイパーパラメータをサンプリング

このノートブックでは、GP ハイパーパラメータをサンプリングするために GPyTorch と NUTS をどのように統合するか、そして完全に Bayesian な方法で GP 推論を遂行するかを実演します。

GPyTorch でのサンプリングの高位概要は以下のようなものです :

  1. ExactGP を拡張して forward メソッドを定義して、通常のように貴方のモデルを定義します。
  2. 貴方のモデルが定義する各パラメータについて、そのパラメータ、あるいはパラメータのある関数で GPyTorch 事前分布を登録する必要があります。デフォルト closure 以外の何かを利用する場合には (e.g., パラメータか変換されたパラメータ名を指定することにより)、setting_closure を指定する必要もあります : gpytorch.Module.register_prior のための docs を見てください。
  3. pyro モデルを定義します、これは各 GP パラメータのためのサンプル site を持ち、それから損失を計算します。便利のために、gpytorch.Module 上で pyro_sample_from_prior メソッドを定義します、これは前者の演算を行ないます。後者の演算については、損失を得るために mll(output, y) の代わりに mll.pyro_factor(output, y) を単に呼び出します。
  4. サンプルを生成するために貴方がちょうど定義した pyro モデル上で NUTS (or HMC etc) を実行します。これは貴方が定義した事前分布に依拠してかなり時間がかかるか全く時間がかからない可能性があることに注意してください。
  5. サンプルをモデルにロードし、モデルを単純な GP からバッチ GP に変換します (単純なバッチ GP 上の例題ノートブック参照)、そこではバッチの 各 GP は異なるハイパーパラメータ・サンプルに対応します。
  6. 各ハイパーパラメータ・サンプルのための予測を得るためにバッチ GP にテストデータを渡します。 
import math
import torch
import gpytorch
import pyro
from pyro.infer.mcmc import NUTS, MCMC
from matplotlib import pyplot as plt

%matplotlib inline
%load_ext autoreload
%autoreload 2

# Training data is 11 points in [0,1] inclusive regularly spaced
train_x = torch.linspace(0, 1, 6)
# True function is sin(2*pi*x) with Gaussian noise
train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2

# We will use the simplest form of GP model, exact inference
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.PeriodicKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

 

サンプリングを実行する

次のセルは他のワークフローとは本質的に異なるコードの最初のピースです。そこでは、モデルと尤度を通常のように作成してから、モデルのパラメータの各々に事前分布を登録します。事前分布を raw パラメータ (e.g., “raw_lengthscale”) ではなく変換されたパラメータ (e.g., “lengthscale”) に直接登録できることに注意してください。これは有用です、けれども事前分布を指定する必要があります、そのサポートはパラメータの領域 (= domain) に完全に含まれます。例えば、lengthscale 事前分布は正の実数かそのサブセットに渡るだけのサポートを持たなければなりません。

# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
num_samples = 2 if smoke_test else 100
warmup_steps = 2 if smoke_test else 200


from gpytorch.priors import LogNormalPrior, NormalPrior, UniformPrior
# Use a positive constraint instead of usual GreaterThan(1e-4) so that LogNormal has support over full range.
likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.Positive())
model = ExactGPModel(train_x, train_y, likelihood)

model.mean_module.register_prior("mean_prior", UniformPrior(-1, 1), "constant")
model.covar_module.base_kernel.register_prior("lengthscale_prior", UniformPrior(0.01, 0.5), "lengthscale")
model.covar_module.base_kernel.register_prior("period_length_prior", UniformPrior(0.05, 2.5), "period_length")
model.covar_module.register_prior("outputscale_prior", UniformPrior(1, 2), "outputscale")
likelihood.register_prior("noise_prior", UniformPrior(0.05, 0.3), "noise")

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

def pyro_model(x, y):
    model.pyro_sample_from_prior()
    output = model(x)
    loss = mll.pyro_factor(output, y)
    return y

nuts_kernel = NUTS(pyro_model, adapt_step_size=True)
mcmc_run = MCMC(nuts_kernel, num_samples=num_samples, warmup_steps=warmup_steps, disable_progbar=smoke_test)
mcmc_run.run(train_x, train_y)

Sample: 100%|█████████████████████████████████████████| 300/300 [00:17, 17.30it/s, step size=6.50e-01, acc. prob=0.895]

 

サンプルをロードする

次のセルでは、NUTS により生成されたサンプルをモデルにロードします。これはモデルを単一の GP から num_samples GP のバッチに変換します、この場合は 100 です。

model.pyro_load_from_samples(mcmc_run.get_samples())

model.eval()
test_x = torch.linspace(0, 1, 101).unsqueeze(-1)
test_y = torch.sin(test_x * (2 * math.pi))
expanded_test_x = test_x.unsqueeze(0).repeat(num_samples, 1, 1)
output = model(expanded_test_x)

 

Mean 関数をプロットする

次のセルでは、同じプロット上で最初の 25 mean 関数をプロットします。この特定の例題は 1 次元のためだけの非常に巨大な総量のデータを持ちますので、ハイパーパラメータ事後分布は非常にタイトで比較的小さい分散があります。

with torch.no_grad():
    # Initialize plot
    f, ax = plt.subplots(1, 1, figsize=(4, 3))

    # Plot training data as black stars
    ax.plot(train_x.numpy(), train_y.numpy(), 'k*', zorder=10)

    for i in range(min(num_samples, 25)):
        # Plot predictive means as blue line
        ax.plot(test_x.numpy(), output.mean[i].detach().numpy(), 'b', linewidth=0.3)

    # Shade between the lower and upper confidence bounds
    # ax.fill_between(test_x.numpy(), lower.numpy(), upper.numpy(), alpha=0.5)
    ax.set_ylim([-3, 3])
    ax.legend(['Observed Data', 'Sampled Means'])

 

ディスクからのモデルのロードをシミュレートする

ディスクから fully Bayesian モデルをロードすることは標準モデルのロードとは僅かに違います、何故ならばサンプリングのプロセスがモデルのパラメータの shape を変更するからです。これに対処するには、state dict をロードする前にモデル上で load_strict_shapes(False) を呼び出す必要があります。下のセルでは、モデルを再作成して state dict からロードすることによりこれを実演します。

load_strict_shapes 呼び出しなくては、これは失敗するであろうことに注意してください。

state_dict = model.state_dict()
model = ExactGPModel(train_x, train_y, likelihood)

# Load parameters without standard shape checking.
model.load_strict_shapes(False)

model.load_state_dict(state_dict)
 

以上