PyTorch : Pyro examples : 変分オートエンコーダ (翻訳)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 11/11/2018 (v0.2.1)
* 本ページは、Pyro のドキュメント Examples : Variational Autoencoders を翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
変分オートエンコーダ
イントロダクション
変分オートエンコーダ (VAE) は間違いなく深層確率的モデリングを実現する最も単純なセットアップです。ここでは術語の選択に私達は慎重でいることに注意してください。VAE はそれ自体はモデルではありません — むしろ VAE はモデルの特定のクラスのために変分推論を行なうための特定のセットアップです。モデルのクラスは非常に幅広いです: 基本的には潜在確率変数を持つ任意の (教師なし) 密度推定器です。そのようなモデルの基本構造は単純で、殆ど見かけによらずそのようなものです (Fig. 1 参照)。
ここで興味がある種類のモデルの構造をグラフィカル・モデルとして描写しました。私達は \(N\) 観測ポイント \(\{ \bf x_i \}\) を持ちます。各データポイントは (局所) 潜在確率変数 \(\bf z_i\) により生成されます。パラメータ \(\theta\) もあります、これは総てのデータポイントがそれに依拠するという意味でグローバルです (これが何故それが矩形の外側に描かれているかです)。\(\theta\) はパラメータですから、それについてベイジアンである何かではないことに注意してください。最後に、ここで特に重要なことは 各 \(\bf x_i\) に複雑で、非線形な方法で \(\bf z_i\) に依拠することを許容することです。実際にこの依存性はパラメータ \(\theta\) を持つ (深層) ニューラルネットワークによりパラメータ化されます。モデルのこのクラスのための推論を特に挑戦的にしているのはこの非線形です。
もちろんこの非線形構造もまた複雑なデータをモデル化するためにこのモデルのクラスが非常に柔軟なアプローチを提供する一つの理由です。実際にモデルの構成要素の各々が様々な異なる方法で「再構成 (= reconfigure)」可能であることは強調するに値します。例えば :
- \(p_\theta({\bf x} | {\bf z})\) のニューラルネットワークは通常の総ての方法で変更されます (層数、非線形のタイプ、隠れユニットの数, etc.)。
- 手元のデータセットに適合する観測尤度を選択できます : gaussian, bernoulli, categorical, etc.
- 潜在空間の次元数を選択できます。
グラフィカルモデル表現はモデルの構造について考えるための有用な方法ですが、同時確率密度の明示的な分解を見ることもまた有益です :
\[
p({\bf x}, {\bf z}) = \prod_{i=1}^N p_\theta({\bf x}_i | {\bf z}_i) p({\bf z}_i)
\]
\(p({\bf x}, {\bf z})\) がこのような項の積に分解されるという事実は \(\bf z_i\) を局所確率変数と呼ぶとき意味するものを明確にします。任意の特定の\(i\) について、単一データポイント \(\bf x_i\) だけが \(\bf z_i\) に依拠します。そのようなものとして \(\{\bf z_i\}\) はローカル構造 i.e. 各データポイントに private な構造を記述します。分解された構造は学習の過程で部分サンプリングを行えることもまた意味します。そのようなものとしてこの種のモデルは巨大なデータ設定に従います (= amenable)。(これと関連トピックの更なる議論については SVI Part II を見てください。)
モデルについてはそれだけのことです。観測が複雑で、非線型な方法で潜在確率変数に依拠しますので、潜在変数に渡る事後分布に複雑な構造を持つことを期待します。その結果このモデルで推論するためにガイド (i.e. 変分分布) の融通の効く族を指定する必要があります。巨大なデータセットにスケールできることを望むので、ガイドは変分パラメータの数を制御下に保持するために amortization を活用していきます (amortization のある程度より一般的な議論は SVI Part II 参照)。
ガイドのジョブは潜在確率変数のために良い値を「推測 (= guess)」することです — それらはモデル事前分布に対して真でありそしてデータに対して真であるという意味で「良い」です。amortization を利用しないのであれば、各データポイント \(\bf x_i\) に対して変分パラメータ \(\{ \lambda_i \}\) を導入します。これらの変分パラメータは \(\bf z_i\) の「良い」値についての信念を表します ; 例えば、それらは \({\bf z}_i\) 空間のガウス分布の平均と分散をエンコードするかもしれません。amortization は変分パラメータ \(\{ \lambda_i \}\) を導入するよりも寧ろ、代わりに各 \(\bf x_i\) を適切な \(\lambda_i\) にマップする関数を学習することを意味します。この関数が柔軟であることを必要としますので、それをニューラルネットワークとしてパラメータ化します。こうして総ての \(N\) データポイント \({\bf x}_i\) に対してインスタン化できる潜在的 \(\bf z\) 空間に渡る分布のパラメータ化された族で終わります (fig. 2 参照)。
ガイド \(q_{\phi}({\bf z} | {\bf x})\) は、総てのデータポイントにより共有される大域パラメータ \(\phi\) によりパラメータ化されることに注意してください。推論の目標は 2 つの条件が満たされるように \(\theta\) と \(\phi\) のための「良い」値を見つけることです :
- log エビデンス \(\log p_\theta({\bf x})\) は巨大です。これはモデルがデータへの良い fit であることを意味します。
- ガイド \(q_{\phi}({\bf z} | {\bf x})\) は事後分布への良い近似を提供します。
(確率的変分推論へのイントロダクションについては SVI Part I 参照。)
この時点でズームアウトして私達のセットアップの高位構造を考えます。具体性のために、\(\{ \bf x_i \}\) を画像としてモデルが画像の生成モデルであると仮定しましょう。ひとたび \(\theta\) の良い値を学習したならば次のようにモデルから画像を生成できます :
- 事前分布 \(p({\bf z})\) に従って \(\bf z\) をサンプリングする
- 尤度 \(p_\theta({\bf x}|{\bf z})\) に従って\(\bf x\) をサンプリングする
各画像は潜在的コード \(\bf z\) により表現されてそのコードは尤度を使用して画像にマップされます、これは 学習した \(\theta\) に依拠します。これがこのコンテキストで尤度がしばしばデコーダと呼ばれる理由です : そのジョブは \(\bf z\) を \(\bf x\) にデコードすることです。これは確率モデルですので、与えられたデータポイント \(\bf x\) をエンコードする \(\bf z\) についての不確かさがあることに注意してください。
ひとたび \(\theta\) と \(\phi\) のための良い値を学習したのであれば次の課題を通り抜けます :
- 与えられた画像 \(\bf x\) で始めます。
- ガイドを使用してそれを \(\bf z\) としてエンコードします。
- モデル尤度を使用して \(\bf z\) をデコードして再構築された画像 \({\bf x}_\rm{reco}\) を得ます。
\(\theta\) と \(\phi\) のための良い値を学習したのであれば、\(\bf x\) と \({\bf x}_\rm{reco}\) は類似するはずです。これは用語オートエンコーダがこのセットアップを記述するために何故使用されることになるかを明らかにするはずです。まとめて、それらはオートエンコーダとして考えることができます。
Pyro の VAE
前置きはこれくらいにしておきましょう。Pyro でどのように VAE を実装するかを見てみましょう。モデル化していくデータセットは MNIST です、手書き数字の画像のコレクションです。これはポピュラーなベンチマーク・データセットですので、書く必要があるボイラープレート・コードの総量を減らすために PyTorch の便宜的なデータローダ機能を利用できます :
import os import numpy as np import torch import torchvision.datasets as dset import torch.nn as nn import torchvision.transforms as transforms import pyro import pyro.distributions as dist from pyro.infer import SVI, Trace_ELBO from pyro.optim import Adam
pyro.enable_validation(True) pyro.distributions.enable_validation(False) pyro.set_rng_seed(0) # Enable smoke test - run the notebook cells on CI. smoke_test = 'CI' in os.environ
# for loading and batching MNIST dataset
def setup_data_loaders(batch_size=128, use_cuda=False):
root = './data'
download = True
trans = transforms.ToTensor()
train_set = dset.MNIST(root=root, train=True, transform=trans,
download=download)
test_set = dset.MNIST(root=root, train=False, transform=trans)
kwargs = {'num_workers': 1, 'pin_memory': use_cuda}
train_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset=test_set,
batch_size=batch_size, shuffle=False, **kwargs)
return train_loader, test_loader
ここで注目を集める主なことはピクセル濃度を範囲 [0.0, 1.0] に正規化するために transforms.ToTensor() を使用することです。
次に私達の Decoder ネットワークをカプセル化する PyTorch モジュールを定義します :
class Decoder(nn.Module):
def __init__(self, z_dim, hidden_dim):
super(Decoder, self).__init__()
# setup the two linear transformations used
self.fc1 = nn.Linear(z_dim, hidden_dim)
self.fc21 = nn.Linear(hidden_dim, 784)
# setup the non-linearities
self.softplus = nn.Softplus()
self.sigmoid = nn.Sigmoid()
def forward(self, z):
# define the forward computation on the latent z
# first compute the hidden units
hidden = self.softplus(self.fc1(z))
# return the parameter for the output Bernoulli
# each is of size batch_size x 784
loc_img = self.sigmoid(self.fc21(hidden))
return loc_img
潜在的変数 z が与えられたとき、Decoder の forward 呼び出しは画像空間の Bernoulli 分布のためのパラメータを返します。各画像はサイズ 28×28=784 なので、loc_img はサイズ batch_size x 784 です。
次に Encoder ネットワークをカプセル化する PyTorch モジュールを定義します :
class Encoder(nn.Module):
def __init__(self, z_dim, hidden_dim):
super(Encoder, self).__init__()
# setup the three linear transformations used
self.fc1 = nn.Linear(784, hidden_dim)
self.fc21 = nn.Linear(hidden_dim, z_dim)
self.fc22 = nn.Linear(hidden_dim, z_dim)
# setup the non-linearities
self.softplus = nn.Softplus()
def forward(self, x):
# define the forward computation on the image x
# first shape the mini-batch to have pixels in the rightmost dimension
x = x.reshape(-1, 784)
# then compute the hidden units
hidden = self.softplus(self.fc1(x))
# then return a mean vector and a (positive) square root covariance
# each of size batch_size x z_dim
z_loc = self.fc21(hidden)
z_scale = torch.exp(self.fc22(hidden))
return z_loc, z_scale
画像 x が与えられたとき Encoder の forward 呼び出しは平均と共分散を返します、それらは一緒に潜在空間の (対角) ガウス分布をパラメータ化します。
手元のエンコーダとデコーダ・ネットワークで、今はモデルとガイドを表わす確率関数を書き下すことができます。最初にモデル :
# define the model p(x|z)p(z)
def model(self, x):
# register PyTorch module `decoder` with Pyro
pyro.module("decoder", self.decoder)
with pyro.iarange("data", x.shape[0]):
# setup hyperparameters for prior p(z)
z_loc = x.new_zeros(torch.Size((x.shape[0], self.z_dim)))
z_scale = x.new_ones(torch.Size((x.shape[0], self.z_dim)))
# sample from prior (value will be sampled by guide when computing the ELBO)
z = pyro.sample("latent", dist.Normal(z_loc, z_scale).independent(1))
# decode the latent code z
loc_img = self.decoder.forward(z)
# score against actual images
pyro.sample("obs", dist.Bernoulli(loc_img).independent(1), obs=x.reshape(-1, 784))
model() は画像 x のミニバッチを入力として取る callable であることに注意してください。これはサイズ batch_size x 784 の torch.Tensor です。
model() の内部で行なう最初のことは Pyro で (前にインスタンス化された) decoder モジュールを登録することです。それに適切な (そして一意な) 名前を与えることに注意してください。pyro.module へのこの呼び出しは Pyro に decoder ネットワーク内の総てのパラメータを知らせます。
次に単なる単位正規ガウス分布である事前分布のためのハイパーパラメータをセットアップします。注意してください :- pyro.iarange を通してミニバッチ (i.e. 最左端次元) のデータの中の独立性を特に明示します。また、潜在変数 z からサンプリングするとき .independent(1) の使用方法に注意してください – サンプルを batch_size = z_dim を持つ単変量正規分布から生成されたものとして取り扱う代わりに、それらを対角共分散を持つ多変量正規分布から生成されたものとして取り扱うことを確実にします。より詳細は Tensor Shapes チュートリアルを参照してください。- 画像のミニバッチ全体を処理していますので、ミニバッチ size に等しい z_loc と z_scale の最左端次元が必要です – GPU 上の場合には、新たに作成された tensor が同じ GPU デバイス上にあることを確実にするために new_zeros と new_ones を使用します。
次に確率変数に一意な Pyro 名 ‘latent’ を与えることを確実にして、事前分布から潜在変数 z をサンプリングします。それから loc_img を返すデコーダ・ネットワークを通して z を渡します。そして loc_img によりパラメータ化された Bernoulli 尤度に対してミニバッチ x の観測された画像をスコアします。総てのピクセルが最右端次元にあるように x を平坦化することに注意してください。
That’s all there is to it! モデルの Pyro プリミティブのフローがモデルの生成ストーリーをどれほど近く追っているかに注意してください、e.g. Figure 1 によるカプセル化されたものとして。さてガイドに進みます :
# define the guide (i.e. variational distribution) q(z|x)
def guide(self, x):
# register PyTorch module `encoder` with Pyro
pyro.module("encoder", self.encoder)
with pyro.iarange("data", x.shape[0]):
# use the encoder to get the parameters used to define q(z|x)
z_loc, z_scale = self.encoder.forward(x)
# sample the latent code z
pyro.sample("latent", dist.Normal(z_loc, z_scale).independent(1))
ちょうどモデルでのように、最初に使用している PyTorch モジュール (つまり encoder) を Pyro で登録します。画像 x のミニバッチを取りそれを encoder を通して渡します。それから encoder ネットワークにより出力されたパラメータを使用してミニバッチの各画像に対して潜在変数の値をサンプリングするために正規分布を使用します。重要なこととして、(モデルでもそうしたように) 潜在確率変数のために同じ名前を使用します : ‘latent’。また、ミニバッチ次元の独立性を明示するための pyro.iarange と、そして z_dims 上の依存性を強制する .independent(1) の利用方法には注意してください、モデルで行なったことと正確に同じです。
完全なモデルとガイドを定義した今、推論に進むことができます。しかしそれを行なう前に PyTorch モジュールでモデルとガイドをどのようにパッケージ化するかを見ましょう :
class VAE(nn.Module):
# by default our latent space is 50-dimensional
# and we use 400 hidden units
def __init__(self, z_dim=50, hidden_dim=400, use_cuda=False):
super(VAE, self).__init__()
# create the encoder and decoder networks
self.encoder = Encoder(z_dim, hidden_dim)
self.decoder = Decoder(z_dim, hidden_dim)
if use_cuda:
# calling cuda() here will put all the parameters of
# the encoder and decoder networks into gpu memory
self.cuda()
self.use_cuda = use_cuda
self.z_dim = z_dim
# define the model p(x|z)p(z)
def model(self, x):
# register PyTorch module `decoder` with Pyro
pyro.module("decoder", self.decoder)
with pyro.iarange("data", x.shape[0]):
# setup hyperparameters for prior p(z)
z_loc = x.new_zeros(torch.Size((x.shape[0], self.z_dim)))
z_scale = x.new_ones(torch.Size((x.shape[0], self.z_dim)))
# sample from prior (value will be sampled by guide when computing the ELBO)
z = pyro.sample("latent", dist.Normal(z_loc, z_scale).independent(1))
# decode the latent code z
loc_img = self.decoder.forward(z)
# score against actual images
pyro.sample("obs", dist.Bernoulli(loc_img).independent(1), obs=x.reshape(-1, 784))
# define the guide (i.e. variational distribution) q(z|x)
def guide(self, x):
# register PyTorch module `encoder` with Pyro
pyro.module("encoder", self.encoder)
with pyro.iarange("data", x.shape[0]):
# use the encoder to get the parameters used to define q(z|x)
z_loc, z_scale = self.encoder.forward(x)
# sample the latent code z
pyro.sample("latent", dist.Normal(z_loc, z_scale).independent(1))
# define a helper function for reconstructing images
def reconstruct_img(self, x):
# encode image x
z_loc, z_scale = self.encoder(x)
# sample in latent space
z = dist.Normal(z_loc, z_scale).sample()
# decode the image (note we don't sample in image space)
loc_img = self.decoder(z)
return loc_img
ここで強調したいことは 2 つのモジュール encoder と decoder は VAE の属性であることです (それ自身は nn.Module から継承しています)。これはそれらが両者とも VAE モジュールに属するものとして自動的に登録されるという結果になります。従って、例えば、VAE のインスタンス上で parameters() を呼び出すとき、PyTorch は総ての関連パラメータを返すことを知ります。それはまた GPU 上で実行する場合、cuda() への呼び出しが総ての (サブ) モジュールの総てのパラメータを GPU メモリに移動することも意味します。
推論
さて推論のための準備ができました。次のセクションで完全なコードに言及します。
最初に VAE モジュールのインスタンスをインスタンス化します。
vae = VAE()
それから Adam optimizer のインスタンスをセットアップします。
optimizer = Adam({"lr": 1.0e-3})
それから推論アルゴリズムをセットアップします、これは ELBO を最大化することによりモデルとガイドのために良いパラメータを学習していきます :
svi = SVI(vae.model, vae.guide, optimizer, loss=Trace_ELBO())
That’s all there is to it. さて訓練ループを定義しなければならないだけです :
def train(svi, train_loader, use_cuda=False):
# initialize loss accumulator
epoch_loss = 0.
# do a training epoch over each mini-batch x returned
# by the data loader
for x, _ in train_loader:
# if on GPU put mini-batch into CUDA memory
if use_cuda:
x = x.cuda()
# do ELBO gradient and accumulate loss
epoch_loss += svi.step(x)
# return epoch loss
normalizer_train = len(train_loader.dataset)
total_epoch_loss_train = epoch_loss / normalizer_train
return total_epoch_loss_train
総てのミニバッチ・ロジックは data loader により処理されることに注意してください。訓練ループの要点は svi.step(x) です。ここで注目すべきことは 2 つあります :
- step への総ての引数はモデルとガイドに渡されます。結果としてモデルとガイドは同じ呼び出しシグネチャを持つ必要があります。
- step は損失のノイズを持つ推定を返します (i.e. minus ELBO)。この推定はどのような方法でも正規化されていません、従って e.g. それはミニバッチのサイズでスケールします。
評価ロジックを追加するためのロジックは類推的です :
def evaluate(svi, test_loader, use_cuda=False):
# initialize loss accumulator
test_loss = 0.
# compute the loss over the entire test set
for x, _ in test_loader:
# if on GPU put mini-batch into CUDA memory
if use_cuda:
x = x.cuda()
# compute ELBO estimate and accumulate loss
test_loss += svi.evaluate_loss(x)
normalizer_test = len(test_loader.dataset)
total_epoch_loss_test = test_loss / normalizer_test
return total_epoch_loss_test
基本的には行なう必要がある唯一の変更は step の代わりに evaluate_loss を呼び出すことです。この関数は ELBO の推定を計算しますが、どのような勾配ステップも取りません。
ハイライトしたいコードの最後のピースは VAE クラスのヘルパー・メソッド reconstruct_img です : これは (コードに翻訳された) イントロダクションで説明した単なる画像再構築実験です。画像を取りそれをエンコーダを通して渡します。それからエンコーダから提供されるガウス分布を使用して潜在空間でサンプリングします。最後に潜在コードを画像にデコードします : それでサンプリングする代わりに平均ベクトル loc_img を返します。sample() ステートメントは確率的ですので、reconstruct_img 関数を実行するたびに z の異なるドローを得ることに注意してください。良いモデルとガイドを学習したのであれば — 特に良い潜在表現を学習したのであれば — z サンプルのこの大多数は書き数字の異なるスタイルに対応するでしょう、そして再構築された画像は興味深い様々な異なるスタイルを展示します。
コードとサンプル結果
訓練は訓練データセットに渡りエビデンス下限 (ELBO, evidence lower bound) を最大化することに相当します。100 反復のために訓練してテストデータセットのために ELBO を評価します。Figure 3 を見てください。
# Run options LEARNING_RATE = 1.0e-3 USE_CUDA = False # Run only for a single iteration for testing NUM_EPOCHS = 1 if smoke_test else 100 TEST_FREQUENCY = 5
train_loader, test_loader = setup_data_loaders(batch_size=256, use_cuda=USE_CUDA)
# clear param store
pyro.clear_param_store()
# setup the VAE
vae = VAE(use_cuda=USE_CUDA)
# setup the optimizer
adam_args = {"lr": LEARNING_RATE}
optimizer = Adam(adam_args)
# setup the inference algorithm
svi = SVI(vae.model, vae.guide, optimizer, loss=Trace_ELBO())
train_elbo = []
test_elbo = []
# training loop
for epoch in range(NUM_EPOCHS):
total_epoch_loss_train = train(svi, train_loader, use_cuda=USE_CUDA)
train_elbo.append(-total_epoch_loss_train)
print("[epoch %03d] average training loss: %.4f" % (epoch, total_epoch_loss_train))
if epoch % TEST_FREQUENCY == 0:
# report test diagnostics
total_epoch_loss_test = evaluate(svi, test_loader, use_cuda=USE_CUDA)
test_elbo.append(-total_epoch_loss_test)
print("[epoch %03d] average test loss: %.4f" % (epoch, total_epoch_loss_test))
次にモデルからランダムにサンプリングされた画像のセットを示します。これらは z のランダム・サンプルをドローして各々一つに対して画像を生成することによりこれらは生成されます、Figure 4 を見てください。
私達はまた総ての MNIST 画像をエンコードしてそれらの平均を 2-次元 T-SNE 空間に埋め込むことによりテストデータセット全体の 50-次元潜在空間を学習します。それからそのクラスにより各々の埋め込まれた画像を彩色します。結果の Figure 5 は各クラス・クラスタ内の分散を持つクラスにより分離を示します。
GitHub 上のフルコードを見てください。
References
- [1] Auto-Encoding Variational Bayes, Diederik P Kingma, Max Welling
- [2] Stochastic Backpropagation and Approximate Inference in Deep Generative Models, Danilo Jimenez Rezende, Shakir Mohamed, Daan Wierstra
以上