pythonの確率的プログラミングのライブラリであるEdwardは元々計算にtensorflowを使っていましたが、発展版のEdward2は TensorFlow Probability の一部として取り込まれました。
クラスや関数が大きく変わり互換性がないので相違点に関する文書が公開されています。
以下の内容は
https://github.com/tensorflow/probability/blob/22cb49d4b7d70a96a46f3a230792eab7ef793cb8/tensorflow_probability/python/edward2/Upgrading_From_Edward_To_Edward2.md
の内容の概略(劣化?)です。コードの部分は同一のものを使っています。
以下ではDeep Exponential Family(Edward, TensorFlow Probability)の例を比較してます。
名前空間
Edward
import edward as ed
from edward.models import Empirical, Gamma, Poisson
Edward2,TFP
import tensorflow_probability as tfp
from tensorflow_probability import edward2 as ed
モデルの書き方
大きな違いはEdwardではObjectであったのに対しEdward2では分布を返す関数になるというところです。
bag_of_words = np.random.poisson(5., size=[256, 32000]) # training data as matrix of counts
data_size, feature_size = bag_of_words.shape # number of documents x words (vocabulary)
units = [100, 30, 15] # number of stochastic units per layer
shape = 0.1 # Gamma shape parameter
w2 = Gamma(0.1, 0.3, sample_shape=[units[2], units[1]])
w1 = Gamma(0.1, 0.3, sample_shape=[units[1], units[0]])
w0 = Gamma(0.1, 0.3, sample_shape=[units[0], feature_size])
z2 = Gamma(0.1, 0.1, sample_shape=[data_size, units[2]])
z1 = Gamma(shape, shape / tf.matmul(z2, w2))
z0 = Gamma(shape, shape / tf.matmul(z1, w1))
x = Poisson(tf.matmul(z1, w0))
def deep_exponential_family(data_size, feature_size, units, shape):
"""A multi-layered topic model over a documents-by-terms matrix."""
w2 = ed.Gamma(0.1, 0.3, sample_shape=[units[2], units[1]], name="w2")
w1 = ed.Gamma(0.1, 0.3, sample_shape=[units[1], units[0]], name="w1")
w0 = ed.Gamma(0.1, 0.3, sample_shape=[units[0], feature_size], name="w0")
z2 = ed.Gamma(0.1, 0.1, sample_shape=[data_size, units[2]], name="z2")
z1 = ed.Gamma(shape, shape / tf.matmul(z2, w2), name="z1")
z0 = ed.Gamma(shape, shape / tf.matmul(z1, w1), name="z0")
x = ed.Poisson(tf.matmul(z0, w0), name="x")
return x
Session
tensorflowのSessionでrunを行うのは同様ですがEdwardでは専用のed.get_session()を使って実行します。
with ed.get_session() as sess:
sess.run(x)
Edward2では関数であるモデルの出力xをrunします。ここではまだデータとの関連付けはされません。
x = deep_exponential_family(data_size, feature_size, units, shape)
with tf.Session() as sess: # or, e.g., tf.train.MonitoredSession()
sess.run(x)
このように他の確率的プログラミングのライブラリと同様モデルの記述と推論を分けるような形なのでtensorflow eagerでできるようなdefine-by-runの書き方とは相容れなさそうと思ったらそうでもなく、tf.enable_eager_execution()を使うことが可能だそうです。x.numpy()とすることでnumpy arrayの値を得ることができます。
推論
Edwardではマルコフ連鎖モンテカルロ法(MCMC),変分推論(VI)それぞれのクラスにモデルを代入してそれぞれのメソッドで推論計算を実行します。Edward2ではMCMC, VIでそれぞれ異なる形で一手間増えている印象です。
変分推論
EdwardではKL divergenceのクラスのコンストラクタにモデルの確率分布、近似分布を辞書で対応づけて入力します。
def trainable_positive_pointmass(shape, name=None):
"""Learnable point mass distribution over positive reals."""
with tf.variable_scope(None, default_name="trainable_positive_pointmass"):
return PointMass(tf.nn.softplus(tf.get_variable("mean", shape)), name=name)
def trainable_gamma(shape, name=None):
"""Learnable Gamma via shape and scale parameterization."""
with tf.variable_scope(None, default_name="trainable_gamma"):
return Gamma(tf.nn.softplus(tf.get_variable("shape", shape)),
1.0 / tf.nn.softplus(tf.get_variable("scale", shape)),
name=name)
qw2 = trainable_positive_pointmass(w2.shape)
qw1 = trainable_positive_pointmass(w1.shape)
qw0 = trainable_positive_pointmass(w0.shape)
qz2 = trainable_gamma(z2.shape)
qz1 = trainable_gamma(z1.shape)
qz0 = trainable_gamma(z0.shape)
inference = ed.KLqp({w0: qw0, w1: qw1, w2: qw2, z0: qz0, z1: qz1, z2: qz2},
data={x: bag_of_words})
Edward2ではed.interceptionという関数を使ってモデルの確率分布、近似分布を対応づけます。
コードの例ではwith構文を使っているがかなりわかりづらいです。with構文の中でモデルと近似分布が対応づけられた予測事後分布を受け取りその後にデータと対応づけます。
deep_exponential_family_variationalという関数でまとめて近似分布を返すようにしています。
def deep_exponential_family_variational():
"""Posterior approx. for deep exponential family p(w{0,1,2}, z{1,2,3} | x)."""
qw2 = trainable_positive_pointmass(w2.shape, name="qw2") # same func as above but with ed2 rv's
qw1 = trainable_positive_pointmass(w1.shape, name="qw1")
qw0 = trainable_positive_pointmass(w0.shape, name="qw0")
qz2 = trainable_gamma(z2.shape, name="qz2") # same func as above but with ed2 rv's
qz1 = trainable_gamma(z1.shape, name="qz1")
qz0 = trainable_gamma(z0.shape, name="qz0")
return qw2, qw1, qw0, qz2, qz1, qz0
def make_value_setter(**model_kwargs):
"""Creates a value-setting interceptor."""
def set_values(f, *args, **kwargs):
"""Sets random variable values to its aligned value."""
name = kwargs.get("name")
if name in model_kwargs:
kwargs["value"] = model_kwargs[name]
return ed.interceptable(f)(*args, **kwargs)
return set_values
# Compute expected log-likelihood. First, sample from the variational
# distribution; second, compute the log-likelihood given the sample.
qw2, qw1, qw0, qz2, qz1, qz0 = deep_exponential_family_variational()
with ed.tape() as model_tape:
with ed.interception(make_value_setter(w2=qw2, w1=qw1, w0=qw0,
z2=qz2, z1=qz1, z0=qz0)):
posterior_predictive = deep_exponential_family(data_size, feature_size, units, shape)
log_likelihood = posterior_predictive.distribution.log_prob(bag_of_words)
# Compute analytic KL-divergence between variational and prior distributions.
kl = 0.
for rv_name, variational_rv in [("z0", qz0), ("z1", qz1), ("z2", qz2),
("w0", qw0), ("w1", qw1), ("w2", qw2)]:
kl += tf.reduce_sum(variational_rv.distribution.kl_divergence(
model_tape[rv_name].distribution))
elbo = tf.reduce_mean(log_likelihood - kl)
tf.summary.scalar("elbo", elbo)
optimizer = tf.train.AdamOptimizer(1e-3)
train_op = optimizer.minimize(-elbo)
ELBO(evidence lower bound )を定義に近い形(事後分布の対数尤度+KL divergence)で書き、それを-lossとしてtensorflowのoptimizerに入れてsession.runで実行するというtensorflowのプログラミングの流れに入れることができます。
これとは別にDenseFlipoutという関数でパラメータの近似分布を含み、lossを返せるようなニューラルネットの層を定義してKL、対数尤度の計算にもちいる方法が定義されているようです
logistic regressionの例、こちらの質問も参照)。
MCMC
EdwardではEmpiricalクラスを使ってVIの近似分布相当の変数を定義します。これはstan, pymcでは必要のないものでした。VIと同様にed.HMCなどMCMCの方法別のクラスのコンストラクタにモデルの確率分布、近似分布を辞書で対応づけて入力します。
num_samples = 10000 # number of events to approximate posterior
qw2 = Empirical(tf.get_variable("qw2/params", [num_samples, units[2], units[1]]))
qw1 = Empirical(tf.get_variable("qw1/params", [num_samples, units[1], units[0]]))
qw0 = Empirical(tf.get_variable("qw0/params", [num_samples, units[0], feature_size]))
qz2 = Empirical(tf.get_variable("qz2/params", [num_samples, data_size, units[2]]))
qz1 = Empirical(tf.get_variable("qz1/params", [num_samples, data_size, units[1]]))
qz0 = Empirical(tf.get_variable("qz0/params", [num_samples, data_size, units[0]]))
inference = ed.HMC({w0: qw0, w1: qw1, w2: qw2, z0: qz0, z1: qz1, z2: qz2},
data={x: bag_of_words})
Edward2でも似たような変数を定義する必要があります。tfp.mcmc.HamiltonianMonteCarloにはed.make_log_joint_fnという関数でモデルから得られる対数分布をデータに対してカリー化したもの(例のtarget_log_prob_fnの出力)を入れて、さらにサンプリングの関数tfp.mcmc.sample_chainにそれを渡します。current_stateとして近似分布相当の変数を入力します。
num_samples = 10000 # number of events to approximate posterior
qw2 = tf.nn.softplus(tf.random_normal([units[2], units[1]])) # initial state
qw1 = tf.nn.softplus(tf.random_normal([units[1], units[0]]))
qw0 = tf.nn.softplus(tf.random_normal([units[0], feature_size]))
qz2 = tf.nn.softplus(tf.random_normal([data_size, units[2]]))
qz1 = tf.nn.softplus(tf.random_normal([data_size, units[1]]))
qz0 = tf.nn.softplus(tf.random_normal([data_size, units[0]]))
log_joint = ed.make_log_joint_fn(deep_exponential_family)
def target_log_prob_fn(w2, w1, w0, z2, z1, z0):
"""Target log-probability as a function of states."""
return log_joint(data_size, feature_size, units, shape,
w2=w2, w1=w1, w0=w0, z2=z2, z1=z1, z0=z0, x=bag_of_words)
hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_log_prob_fn,
step_size=0.01,
num_leapfrog_steps=5)
states, kernels_results = tfp.mcmc.sample_chain(
num_results=num_samples,
current_state=[qw2, qw1, qw0, qz2, qz1, qz0],
kernel=hmc_kernel,
num_burnin_steps=1000)
pymc4ではこのめんどくさい処理をModelクラスとsample関数の実装として隠蔽しています。
https://github.com/pymc-devs/pymc4/blob/master/pymc4/model/base.py#L34
https://yukinagae.hatenablog.com/entry/2018/09/26/084113
ループ
EdwardではKLやMCMCのオブジェクト(例ではinference)をinitialize()で初期化した後メソッドrun()やupdate()を実行します。
inference.initialize(n_iter=10000)
tf.global_variables_initializer().run()
for _ in range(inference.n_iter):
info_dict = inference.update()
inference.print_progress(info_dict)
inference.finalize()
Edward2ではestimatorの中にある関数を使うとinference.run()と同等のことができるらしいです。例では普通のtensorflowプログラムと同様にtrain_op, elboをrunで実行しています。実行中に計算結果を取得できるのでMCMCのkernel_results.is_accepted関数を使ってacceptance ratioを見ることができるそうです。
max_steps = 10000 # number of training iterations
model_dir = None # directory for model checkpoints
sess = tf.Session()
summary = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(model_dir, sess.graph)
start_time = time.time()
sess.run(tf.global_variables_initializer())
for step in range(max_steps):
start_time = time.time()
_, elbo_value = sess.run([train_op, elbo])
if step % 500 == 0:
duration = time.time() - start_time
print("Step: {:>3d} Loss: {:.3f} ({:.3f} sec)".format(
step, elbo_value, duration))
summary_str = sess.run(summary)
summary_writer.add_summary(summary_str, step)
summary_writer.flush()
Criticism(評価)
モデルの性能評価はEdwardではed.copyで予測事後分布をコピーしたのちed.evaluateという関数に指標の種類を文字列(!)として入れて得ていました。ed.ppc関数で予測事後分布からサンプリングを行い統計指標を計算することができます。
# Build posterior predictive: it is parameterized by a variational posterior sample.
posterior_predictive = ed.copy(
x, {w0: qw0, w1: qw1, w2: qw2, z0: qz0, z1: qz1, z2: qz2})
# Evaluate average log-likelihood of data.
ed.evaluate('log_likelihood', data={posterior_predictive: bag_of_words})
## np.ndarray of shape ()
# Compare TF-IDF on real vs generated data.
def tfidf(bag_of_words):
"""Computes term-frequency inverse-document-frequency."""
num_documents = bag_of_words.shape[0]
idf = tf.log(num_documents) - tf.log(tf.count_nonzero(bag_of_words, axis=0))
return bag_of_words * idf
observed_statistics, replicated_statistics = ed.ppc(
lambda data, latent_vars: tf_idx(data[posterior_predictive]),
{posterior_predictive: bag_of_words},
n_samples=100)
Edward2ではtensorflowの関数を使って同様の計算をする方針のようです。tf.metricsのような関数を用いることもできるそうです。
# See posterior_predictive built in Variational Inference section.
log_likelihood = tf.reduce_mean(posterior_predictive.log_prob(bag_of_words))
## tf.Tensor of shape ()
# Simple version: Compare statistics by sampling from model in a for loop.
observed_statistic = sess.run(tfidf(bag_of_words))
replicated_statistic = tfidf(posterior_predictive)
replicated_statistics = [sess.run(replicated_statistic) for _ in range(100)]
Bayesian neural networks(BNN)ではtf.metrics.accuracyで訓練後の性能を計り、Eight Schoolsという統計モデリングでよく使われるサンプルデータセットでは予測事後分布の結果をプロットすることで評価しています。
感想
pytorchベースのpyroの方がわかりやすい…
サンプルコード
公開されているEdward, Edward2, TensorFlow Probability(以下TFP)のサンプルコード
- Probabilistic PCA (Edward, TensorFlow Probability)
- Eight schools (Edward, TensorFlow Probability)
- Linear mixed effects models (Edward, TensorFlow Probability)
- Variational Autoencoder (Edward, TensorFlow Probability)
- Deep Exponential Family (Edward, TensorFlow Probability)
- Mixture of Gaussians (Edward, TensorFlow Probability)
- Logistic regression (Edward, TensorFlow Probability)