PyTorch (11) Variational Autoencoder

今回は、Variational Autoencoder (VAE) の実験をしてみよう。

実は自分が始めてDeep Learningに興味を持ったのがこのVAEなのだ！VAEの潜在空間をいじって多様な顔画像を生成するデモ（Morphing Faces）を見て、これを音声合成の声質生成に使いたいと思ったのが興味のきっかけだった。

今回の実験は、PyTorchの公式にあるVAEのスクリプトを自分なりに読み解いてまとめてみた結果になっている。

180221-variational-autoencoder.ipynb - Google ドライブ

さっそく実験！いつものimport。

import os
import numpy as np
import torch
import torch.nn as nn
import torch.utils.data
import torch.optim as optim
from torch.autograd import Variable
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image

batch_size = 128
num_epochs = 100
seed = 1
out_dir = './vae_2'

出力先のディレクトリを vae_2 としている。潜在空間が2次元であるVAEの結果を意味する。

cuda = torch.cuda.is_available()
if cuda:
    print('cuda is available!')

if not os.path.exists(out_dir):
    os.mkdir(out_dir)

torch.manual_seed(seed)
if cuda:
    torch.cuda.manual_seed(seed)

MNISTのデータをロードする。前回と同じだけど標準化は [0, 1] にしている。

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data',
                   train=True,
                   download=True,
                   transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=True
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('data',
                   train=False,
                   transform=transforms.ToTensor()),
    batch_size=batch_size,
    shuffle=True
)

VAEのアーキテクチャ

VAEは、Autoencoderと似ているが、Encoderの出力が正規分布の平均と共分散行列になり、潜在表現zがその正規分布からサンプリングされる点が異なる。潜在表現zはランダムサンプリングされるため同じ入力画像Xを入れても毎回異なるzにマッピングされる。

f:id:aidiary:20180228212323p:plain

上図のようにEncoderは正規分布の平均ベクトル（mu）と分散共分散行列の対数（logvar）を出力する。ここでは、潜在表現 z は可視化しやすいように2次元としたので、平均ベクトルのサイズは2となる。分散共分散行列は対角行列と仮定しているため対角成分のみ取ってこちらも2次元ベクトルとなる。対数を取る理由がはっきりしないけどアンダーフローを防ぐためかな？

Encoderが出力した平均と分散を持つ正規分布から入力Xの潜在表現zをサンプリングする。

$z \sim N(\mu(X), \Sigma(X))$

ただ、これを普通にやると誤差逆伝搬ができないので Reparameterization Trick というのを使う。上の式でサンプリングするのではなく、

$\epsilon \sim N(0, I)$

$z = \mu(X) + \Sigma(X)^{\frac{1}{2}} \epsilon$

でzを計算する。このように計算すると和と積の演算だけで構成されるので計算グラフが構築でき、誤差逆伝搬ができるとのこと。PyTorchには normal_(mean=0, std=1) という正規乱数を生成するTensor Operationが実装されている。

ここがちょっとわからない。meanとstdを指定しても乱数生成できるみたいだけどReparameterization Trick必要なのかな？多次元正規分布になるとできないのだろうか？

上のアーキテクチャをPyTorchのコードで書くと下のようになる。

class VAE(nn.Module):

    def __init__(self):
        super(VAE, self).__init__()
        
        self.fc1 = nn.Linear(784, 512)
        self.fc21 = nn.Linear(512, 2)  # mu
        self.fc22 = nn.Linear(512, 2)  # logvar

        self.fc3 = nn.Linear(2, 512)
        self.fc4 = nn.Linear(512, 784)
        
        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()
    
    def encode(self, x):
        h = self.relu(self.fc1(x))
        return self.fc21(h), self.fc22(h)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.data.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h = self.relu(self.fc3(z))
        return self.sigmoid(self.fc4(h))
    
    def forward(self, x):
        x = x.view(-1, 784)
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

model = VAE()
if cuda:
    model.cuda()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

ちょっと数式とコードの対応関係を補足すると logvar は $\log \Sigma$ にあたる。

std = logvar.mul(0.5).exp_() は、

$std = \exp(0.5 * \log \Sigma) = \exp(\log \Sigma ^ \frac{1}{2}) = \Sigma^{\frac{1}{2}}$

となる。よって、eps.mul(std).add_(mu) は、

$\epsilon * \Sigma^{\frac{1}{2}} + \mu$

となり、先の式とコードが一致することがわかる！数式だとεはスカラーっぽいけど実装ではベクトルになっている。潜在表現の各次元ごとに異なる乱数をかけるようだ。

スカラーでもよいのかな？

VAEの損失関数

VAEの損失関数は既存のものではなく、独自の定義が必要。

def loss_function(recon_x, x, mu, logvar):
    # size_average=Falseなのでバッチ内のサンプルの合計lossを求める
    # reconstruction loss 入力画像をどのくらい正確に復元できたか？
    # 数式では対数尤度の最大化だが交差エントロピーlossの最小化と等価
    recon = F.binary_cross_entropy(recon_x, x.view(-1, 784), size_average=False)

    # 潜在空間zに対する正則化項
    # P(z|x) が N(0, I)に近くなる（KL-distanceが小さくなる）ようにする
    kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    return recon + kld

このコードを理解するのがとても難しかった。数式の導出は

を見ていただくほうがよいと思う。最終的には下の目的関数が出て来るのだがそこまでのたどり着き方にはいくつかアプローチがあるようだ。

VAEの最終的な目的関数は、

$\log P(X) - D_{KL} [ Q(z|X) || P(z|X) ] = E [\log P(X|z) ] - D_{KL} [Q(z|X)||P(z)]$

となる。左辺のKL divergenceは $D_{KL} [ Q(z|X) || P(z|X) ] \ge 0$ なので

$\log P(X) \ge E [\log P(X|z) ] - D_{KL} [Q(z|X)||P(z)]$

が成り立つ。たとえば、12 - 2 = 10 のとき 12 >= 10。

左辺がデータXの対数尤度なので生成モデルにおいて最大化したい値になる。右辺は 変分下限（ELBO: evidence lower bound） と呼び、対数尤度の下限となる。ここで、対数尤度を最大化する問題を変分下限を最大化する問題に置き換えるのがポイント。変分下限をできるだけ大きくしてやれば、それより大きい対数尤度も大きくなるというわけ。変分下限を最大化するには、

$E_{Q(z|X)} [\log P(X|z) ] \to \max$

$D_{KL} [Q(z|X) || P(z) ] \to \min$

とすればよい。簡単のため期待値の分布は省略してたけど最初の式をきちんと導出していくと Q(z|X) になることがわかる。

Reconstruction Loss

実はここの理解がちょっと怪しい。尤度最大化って交差エントロピーの最小化と等価で正しい？ときどき一部の論文でlossの式に尤度が書いてあってちょっと混乱する。

$E_{Q(z|X)} [\log P(X|z) ] \to \max$

から見ていく。確率分布で書かれているのでわかりにくいが、Q(z|X) は入力画像Xを潜在空間zにマッピングしているためEncoderとみなせる。また、P(X|z) は潜在空間zから元の画像XにマッピングしているためDecoderとみなせる。つまり、この式は入力画像Xを潜在空間zに落としてそこからXに戻したときの対数尤度を最大化しろという意味だと解釈できる。

この対数尤度の最大化は入力画像 $X$ と再構成画像 $\hat{X}$ の交差エントロピーの最小化とみなせる。つまり、数式で表すと下の部分と一致する。これをReconstruction Lossと呼ぶ。

recon = F.binary_cross_entropy(recon_x, x.view(-1, 784), size_average=False)

KL Divergence

次はこの項。

$D_{KL} [Q(z|X) || P(z) ] \to \min$

先に述べたように Q(z|X) は入力画像Xを潜在空間zにマッピングしているためEncoderとみなせる。つまり、Encoderで入力画像Xをマッピングした分布 Q(z|X) が P(z) に近くなるようにしろという制約 と解釈できる。

ここで、P(z) は簡単のため平均0、分散1の多次元正規分布 N(0, I) と仮定する。

先に書いたようにEncoderの出力は N(μ(X), Σ(X)) の正規分布に従うようにしたため、上の式は

$D_{KL} [N(\mu(X), \Sigma (X)) || N(0, I) ]$

となる。多次元正規分布間のKL Divergenceは下の簡単な式で求まる。

$\displaystyle D_{KL} [N(\mu(X), \Sigma(X)) || N(0, I) ] = - \frac{1}{2} \sum_{k} (1 + \log \Sigma(X) - \mu(X)^2 - \Sigma(X))$

参考: 多変量正規分布の場合のKullback Leibler Divergenceの導出 - Qiita

これをコードで書くと

kld = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

となる。この kld は潜在空間zが N(0, I) にちらばるようにする正則化項とみなせる。あとで実際に潜在空間を描画してみるとこの正則化が正しく機能していることがわかる。

訓練ループ

あとはいつもの訓練ループなので比較的簡単。

def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, (data, _) in enumerate(train_loader):
        if cuda:
            data = Variable(data.cuda())
        else:
            data = Variable(data)
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        train_loss += loss.data[0]
        optimizer.step()
    
    # loss_function() は平均ではなく全サンプルの合計lossを返すのでサンプル数で割る
    train_loss /= len(train_loader.dataset)

    return train_loss    
    

def test(epoch):
    model.eval()
    test_loss = 0
    for batch_idx, (data, _) in enumerate(test_loader):
        if cuda:
            data = Variable(data.cuda(), volatile=True)
        else:
            data = Variable(data, volatile=True)
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        test_loss += loss.data[0]
        
        if epoch % 10 == 0:
            # 10エポックごとに最初のminibatchの入力画像と復元画像を保存
            if batch_idx == 0:
                n = 8
                comparison = torch.cat([data[:n],
                                        recon_batch.view(batch_size, 1, 28, 28)[:n]])
                save_image(comparison.data.cpu(),
                           '{}/reconstruction_{}.png'.format(out_dir, epoch), nrow=n)

    test_loss /= len(test_loader.dataset)

    return test_loss

loss_list = []
test_loss_list = []
for epoch in range(1, num_epochs + 1):
    loss = train(epoch)
    test_loss = test(epoch)

    print('epoch [{}/{}], loss: {:.4f} test_loss: {:.4f}'.format(
        epoch + 1,
        num_epochs,
        loss,
        test_loss))

    # logging
    loss_list.append(loss)
    test_loss_list.append(test_loss)

# save the training model
np.save('loss_list.npy', np.array(loss_list))
np.save('test_loss_list.npy', np.array(test_loss_list))
torch.save(model.state_dict(), 'vae.pth')

学習曲線

学習曲線を描いてみよう。

import matplotlib.pyplot as plt
%matplotlib inline
loss_list = np.load('{}/loss_list.npy'.format(out_dir))
plt.plot(loss_list)
plt.xlabel('epoch')
plt.ylabel('loss')
plt.grid()

f:id:aidiary:20180228225157p:plain

test_loss_list = np.load('{}/test_loss_list.npy'.format(out_dir))
plt.plot(test_loss_list)
plt.xlabel('epoch')
plt.ylabel('test loss')
plt.grid()

f:id:aidiary:20180228225210p:plain

テストlossも減っており、学習が進んでいることが確認できる。

入力画像と再構成画像の比較

上が入力画像で下が再構成画像。

from IPython.display import Image
Image('vae_2/reconstruction_10.png')

f:id:aidiary:20180228225258p:plain

Image('vae_2/reconstruction_100.png')

f:id:aidiary:20180228225314p:plain

10エポック目だとあまり再現できてないが、100エポック目だとある程度は再現できている。ただ、潜在空間を2次元とかなり絞ったのであまりくっきりと再現できていない。潜在空間を20次元にすると

f:id:aidiary:20180228225426p:plain

こんな感じでほぼ再構成できることがわかる。潜在空間が2次元だとテストlossは150くらいで収束するが、20次元にするとテストlossは90台まで下がる。

潜在空間の可視化

最後にテストデータを使って潜在空間を可視化しよう。

model.load_state_dict(torch.load('{}/vae.pth'.format(out_dir),
                                 map_location=lambda storage,
                                 loc: storage))
test_dataset = datasets.MNIST('./data', download=True, train=False, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=10000, shuffle=False)
images, labels = iter(test_loader).next()
images = images.view(10000, -1)

# 784次元ベクトルを2次元ベクトルにencode
z = model.encode(Variable(images, volatile=True))
mu, logvar = z
mu, logvar = mu.data.numpy(), logvar.data.numpy()
print(mu.shape, logvar.shape)

各入力画像Xごとに正規分布のパラメータ μ(X), Σ(X) が出力されるがここでは平均 μ(X) の場所に点をプロットすることにする。ランダムサンプリングすると平均の場所にマッピングされる可能性が一番高い。

import pylab
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 10))
plt.scatter(mu[:, 0], mu[:, 1], marker='.', c=labels.numpy(), cmap=pylab.cm.jet)
plt.colorbar()
plt.xlim((-6, 6))
plt.ylim((-6, 6))
plt.grid()