全体のコードと実行結果 - pymc_vae.ipynb

前回は少しだけ、VAEの話をしましたが、PyMC3というライブラリを使って遊んでみました。思いっきり自作のtheano系ライブラリを使っていますが、本家の例ではkerasを使っており、モデルのtheano.shared変数を全部集めてpymc3.fitに渡す方法があれば、どんなtheano系ライブラリでも活用できると思います。

備忘録として、簡単にPyMC3の解説をすると、

# building model
n_latent = 2
n_batch = 64

vae = ConvVAE(n_latent)  # theano で作ったNN 
xs = tt.tensor4("xs")
xs.tag.test_value = numpy.zeros((n_batch, 1, 28, 28)).astype('float32')

with pm.Model() as model:
    zs = pm.Normal("zs", mu=0, sd=1, shape=(n_batch, n_latent),
                   dtype=theano.config.floatX, total_size=len(data))
    xs_ = pm.Normal("xs_", mu=vae.decode(zs), sd=0.1, observed=xs,
                    dtype=theano.config.floatX, total_size=len(data))

dataはMNISTの学習セット全てが入った多次元配列(データ数, 1, 28, 28)です。あとは前回のVAEにでてくる$x, z$をそのままxs, zsとしています。pm.Model()の中で有向グラフのような、確率変数がどういう分布(今回は$\text{xs} \sim \text{Normal}(\mu, \text{sd})$)から生成されているのかを記述します。注意点としては観測変数にはxs.tag.test_valueとして適当な大きさの入力を代入しておかないと、PyMC3内部のアサーションに引っかかってコンパイルできません。

# fitting model
mean, stddev = vae.encode(xs)
local_RVs = OrderedDict({zs: (mean, stddev)})  # encoded stochastic variable q(z|x)
xs_minibatch = pm.Minibatch(data, n_batch)

with model:
    approx = pm.fit(10000, local_rv=local_RVs,
                    obj_optimizer=pm.adam(learning_rate=1e-4),
                    more_obj_params=list(vae.get_params()),
                    more_replacements={xs: xs_minibatch})

ここで RV とは Random Variable (確率変数) のことだそうです。scikit-learnなどによくある fit 関数と違ってユーザが設定する損失などはなく、pm.fit()は変分ベイズ法なのでobserved=xsとした実際の観測データxsに対する前回導出した周辺尤度の下限(i.e., ELBO, 変分下限)を最大化するようにpm.floatXで設定したパラメタやtheano.shared変数を更新します。今回はニューラルネットを含んでいるので単純な勾配法(+Adamによる更新則)で最適化していますが、いくつかの分布ではもっと良い手法も使えるようです。

ただこの感じだと、前回導出したような事前分布のKLダイバージェンスを分割できているのかとか、そもそもサンプリングの回数とか指定できるのか(たぶん pm.sample() を併用するんじゃないでしょうか？)。まだわかりません...。

次回に続くといいですね...(?)