tensorflow2.0でVariationalオートエンコーダーを作る（記録用）

はじめに
全結合層のみのVariationalオートエンコーダー
- ネットワーク
- 結果
畳み込み層ありのVariationalオートエンコーダー
- ネットワーク
- 結果
つまづいたところ
参考

はじめに

tensorflow2.0でVariationalオートエンコーダーを作ってみました。
全結合層のみのモデルと、畳み込み層ありのモデルを二つを作成して比較しました。
勉強の整理用なので、結果が必要な方は参考にしないでください（あんまりうまくいかなかったので....）。
tensorflow2.0のVAEの良いサンプルが知りたい方は、TFの公式（Convolutional Variational Autoencoder | TensorFlow Core）をご確認ください。

今回作成したコードは以下にあります。
github.com

全結合層のみのVariationalオートエンコーダー

まずは前回（tensorflow2.0でシンプルなオートエンコーダーを作る - たそらぼ）作成したオートエンコーダーを改造して全結合のものを作成しました。

ネットワーク

下図のモデルを訓練しました。層の厚さなどのパラメータはチューニングしていません。

f:id:tasotasoso:20191006013607p:plain — 作成したモデル

class Encoder(layers.Layer):
    def __init__(self):
        super(Encoder, self).__init__()
        self.d1 = Dense(units=64, activation='relu')
        self.d2 = Dense(units=64)
        self.d3 = Dense(units=64)
    def call(self, x):
        x = self.d1(x)
        mean = self.d2(x)
        logvar = self.d3(x)
        return mean, logvar
    
class ReparameterizationTrick(layers.Layer):
    def __init__(self):
        super(ReparameterizationTrick, self).__init__()
    def call(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        z = eps * tf.exp(logvar* .5) + mean
        return z

class Decoder(layers.Layer):
    def __init__(self):
        super(Decoder, self).__init__()
        self.d4 = Dense(units=64, activation='relu')
        self.d5 = Dense(units=784)
    def call(self, z):
        x = self.d4(z)
        x = self.d5(x)
        return x

class Autoencorder(Model):
    def __init__(self):
        super(Autoencorder, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        self.reparameterizationtrick = ReparameterizationTrick()
        
    def call(self, x):        
        mean, logvar = self.encoder(x)
        z = self.reparameterizationtrick(mean, logvar)
        reconstructed = self.decoder(z)
        return reconstructed

model = Autoencorder()

結果

100エポック流して、ELBOは-96くらいでした。

f:id:tasotasoso:20191006001129p:plain — 学習前（全結合）

f:id:tasotasoso:20191006001221p:plain — 学習後（全結合）

なんとなく４の形はみえているんですが、ぼやけてしまっています。

畳み込み層ありのVariationalオートエンコーダー

やっぱり畳み込みの方がいいのかなということで、TFの公式（Convolutional Variational Autoencoder | TensorFlow Core）を参考に、畳み込み層をモデルに入れてみました。

ネットワーク

下図のモデルを訓練しました。層の厚さはTFの公式を参考にしています。TFの公式では分散と平均は同じ出力を半分に割る形になっていますが、作成したモデルは分散と平均で別々に重み行列を持っている点で違いがあります。

class Encoder(layers.Layer):
    def __init__(self):
        super(Encoder, self).__init__()
        self.c1 = Conv2D(filters=32, kernel_size=3, strides=(2, 2), activation='relu')
        self.c2 = Conv2D(filters=64, kernel_size=3, strides=(2, 2), activation='relu')
        self.f = Flatten()
        self.d1 = Dense(units=50)
        self.d2 = Dense(units=50)
    def call(self, x):
        x = self.c1(x)
        x = self.c2(x)
        x = self.f(x)
        mean = self.d1(x)
        logvar = self.d2(x)
        return mean, logvar
    
class ReparameterizationTrick(layers.Layer):
    def __init__(self):
        super(ReparameterizationTrick, self).__init__()
    def call(self, mean, logvar):
        eps = tf.random.normal(shape=mean.shape)
        z = eps * tf.exp(logvar* .5) + mean
        return z

class Decoder(layers.Layer):
    def __init__(self):
        super(Decoder, self).__init__()
        self.d3 = Dense(units=7*7*32, activation='relu')
        self.r = Reshape(target_shape=(7, 7, 32))
        self.c3 = Conv2DTranspose(
              filters=64,
              kernel_size=3,
              strides=(2, 2),
              padding="SAME",
              activation='relu')
        self.c4 = Conv2DTranspose(
              filters=32,
              kernel_size=3,
              strides=(2, 2),
              padding="SAME",
              activation='relu')
        self.c5 = Conv2DTranspose(filters=1, kernel_size=3, strides=(1, 1), padding="SAME")
    def call(self, z):
        x = self.d3(z)
        x = self.r(x)
        x = self.c3(x)
        x = self.c4(x)
        x = self.c5(x)
        return x

class Autoencorder(Model):
    def __init__(self):
        super(Autoencorder, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        self.reparameterizationtrick = ReparameterizationTrick()
        
    def call(self, x):        
        mean, logvar = self.encoder(x)
        z = self.reparameterizationtrick(mean, logvar)
        reconstructed = self.decoder(z)
        return reconstructed

model = Autoencorder()

結果

50エポック流して、ELBOは-95くらいでした。

f:id:tasotasoso:20191006004259p:plain — 学習前（畳み込み層あり）

f:id:tasotasoso:20191006004348p:plain — 学習後（畳み込み層あり）

なんかまだぼやけてますが、全結合層のみのモデの半分の学習でも、かなりくっきり見えるようになってきました。ただ、100エポックに増やしてもぼやけが改善されなかったので、結果としてはイマイチです。

つまづいたところ

損失関数

ELBOを最大化しています。
$\log p(x) \geq ELBO = \mathbb{E}_{q(z|x)} \left[ \log\frac{p(x,z)}{q(z|x)} \right]$
ELBOを最大化することで、デコーダーの対数尤度の最大に近づけている理解です（違ったらすみません）。

def log_normal_pdf(sample, mean, logvar, raxis=1):
    log2pi = tf.math.log(2. * np.pi)
    return tf.reduce_sum(
      -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),
      axis=raxis)

def compute_loss(model, x):
    mean, logvar = model.encoder(x)
    z = model.reparameterizationtrick(mean, logvar)
    x_logit = model.decoder(z)
 
    cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)
    logpx_z = -tf.reduce_sum(cross_ent, axis=[1,2,3])
    logpz = log_normal_pdf(z, 0., 0.)
    logqz_x = log_normal_pdf(z, mean, logvar)
    return -tf.reduce_mean(logpx_z + logpz - logqz_x)

TFの公式に倣い、モンテカルロ積分（モンテカルロ積分 - 人工知能に関する断創録）で期待値を計算しています。
また、 $\log p(x|z)$ は元論文のC.1から、sigmoid_cross_entropyで計算しています。

分散の推定

デコーダーはzを推定するための分散を出力しますが、 $\log \sigma ^2$ が推定されます。 $\sigma$ として使う場合はexpで戻してやる必要があります。

感想

ということで、あんまりうまく行きませんでした。
tensorflow公式チュートリアルのモデルと多少形が違うのですが、そんなに効いてくるかなぁという程度の違いなので、もしかすると実装が間違っているかもです。
ELBOの計算は、デコーダーの出力が $\log \sigma ^2$ であることが分かると理解が深まるので、デコーダーから何が出てくるのか認識することが重要でした。