このページは Cloud Translation API によって翻訳されました。
Switch to English

畳み込み変分オートエンコーダ

TensorFlow.orgで表示 GoogleColabで実行 GitHubでソースを表示ノートブックをダウンロード

このノートブックは変オートエンコーダ(VAE)(訓練方法を示しています12 )。 MNISTデータセットで。 VAEは、オートエンコーダーの確率論的解釈であり、高次元の入力データを取得してそれをより小さな表現に圧縮するモデルです。入力を潜在ベクトルにマッピングする従来のオートエンコーダーとは異なり、VAEは入力データをガウス分布の平均や分散などの確率分布のパラメーターにマッピングします。このアプローチは、画像生成に役立つ連続的で構造化された潜在空間を生成します。

CVAE画像潜在空間

セットアップ

pip install -q tensorflow-probability

# to generate gifs
pip install -q imageio
pip install -q git+https://github.com/tensorflow/docs
WARNING: You are using pip version 20.2.2; however, version 20.2.3 is available.
You should consider upgrading via the '/tmpfs/src/tf_docs_env/bin/python -m pip install --upgrade pip' command.
WARNING: You are using pip version 20.2.2; however, version 20.2.3 is available.
You should consider upgrading via the '/tmpfs/src/tf_docs_env/bin/python -m pip install --upgrade pip' command.
WARNING: You are using pip version 20.2.2; however, version 20.2.3 is available.
You should consider upgrading via the '/tmpfs/src/tf_docs_env/bin/python -m pip install --upgrade pip' command.

from IPython import display

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow_probability as tfp
import time

MNISTデータセットをロードする

各MNIST画像は元々784個の整数のベクトルであり、それぞれが0〜255であり、ピクセルの強度を表します。モデル内のベルヌーイ分布を使用して各ピクセルをモデル化し、データセットを静的に2値化します。

(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()
def preprocess_images(images):
  images = images.reshape((images.shape[0], 28, 28, 1)) / 255.
  return np.where(images > .5, 1.0, 0.0).astype('float32')

train_images = preprocess_images(train_images)
test_images = preprocess_images(test_images)
train_size = 60000
batch_size = 32
test_size = 10000

tf.dataを使用して、データをバッチ処理およびシャッフルします

train_dataset = (tf.data.Dataset.from_tensor_slices(train_images)
                 .shuffle(train_size).batch(batch_size))
test_dataset = (tf.data.Dataset.from_tensor_slices(test_images)
                .shuffle(test_size).batch(batch_size))

tf.keras.Sequentialを使用してエンコーダーとデコーダーのネットワークを定義します

VAEの例では、エンコーダーネットワークとデコーダーネットワークに2つの小さなConvNetを使用しています。文献では、これらのネットワークは、それぞれ推論/認識モデルおよび生成モデルとも呼ばれます。実装を簡素化するためにtf.keras.Sequentialを使用します。以下の説明では、$ x $と$ z $がそれぞれ観測変数と潜在変数を示しているとします。

エンコーダネットワーク

これは、近似事後分布$ q(z | x)$を定義します。これは、観測値を入力として受け取り、潜在表現$ z $の条件付き分布を指定するためのパラメーターのセットを出力します。この例では、分布を対角ガウス分布としてモデル化するだけで、ネットワークは因数分解されたガウス分布の平均パラメーターと対数分散パラメーターを出力します。数値安定性のために、分散の代わりに対数分散を直接出力します。

デコーダーネットワーク

これは、観測値$ p(x | z)$の条件付き分布を定義します。これは、潜在サンプル$ z $を入力として受け取り、観測値の条件付き分布のパラメーターを出力します。 $ p(z)$の前の潜在分布を単位ガウスとしてモデル化します。

再パラメータ化のトリック

トレーニング中にデコーダーのサンプル$ z $を生成するために、入力観測値$ x $が与えられた場合に、エンコーダーによって出力されたパラメーターによって定義された潜在分布からサンプリングできます。ただし、このサンプリング操作では、バックプロパゲーションがランダムノードを通過できないため、ボトルネックが発生します。

これに対処するために、再パラメータ化のトリックを使用します。この例では、次のように、デコーダーパラメーターと別のパラメーター$ \ epsilon $を使用して$ z $を概算します。

$$z = \mu + \sigma \odot \epsilon$$

ここで、$ \ mu $と$ \ sigma $は、それぞれガウス分布の平均と標準偏差を表します。それらはデコーダー出力から導き出すことができます。 $ \ epsilon $は、$ z $の確率を維持するために使用されるランダムノイズと考えることができます。標準正規分布から$ \ epsilon $を生成します。

潜在変数$ z $は、$ \ mu $、$ \ sigma $、および$ \ epsilon $の関数によって生成されるようになりました。これにより、モデルは、エンコーダーの勾配をそれぞれ$ \ mu $および$ \ sigma $を介して逆伝播できます。 、$ \ epsilon $を通じて確率を維持しながら。

ネットワークアーキテクチャ

エンコーダネットワークでは、2つの畳み込み層とそれに続く完全に接続された層を使用します。デコーダーネットワークでは、完全に接続されたレイヤーと、それに続く3つの畳み込み転置レイヤー(一部のコンテキストではデコンボリューションレイヤー)を使用して、このアーキテクチャをミラーリングします。ミニバッチの使用による追加の確率論は、サンプリングによる確率論に加えて不安定性を悪化させる可能性があるため、VAEをトレーニングするときにバッチ正規化の使用を避けるのが一般的な方法であることに注意してください。

class CVAE(tf.keras.Model):
  """Convolutional variational autoencoder."""

  def __init__(self, latent_dim):
    super(CVAE, self).__init__()
    self.latent_dim = latent_dim
    self.encoder = tf.keras.Sequential(
        [
            tf.keras.layers.InputLayer(input_shape=(28, 28, 1)),
            tf.keras.layers.Conv2D(
                filters=32, kernel_size=3, strides=(2, 2), activation='relu'),
            tf.keras.layers.Conv2D(
                filters=64, kernel_size=3, strides=(2, 2), activation='relu'),
            tf.keras.layers.Flatten(),
            # No activation
            tf.keras.layers.Dense(latent_dim + latent_dim),
        ]
    )

    self.decoder = tf.keras.Sequential(
        [
            tf.keras.layers.InputLayer(input_shape=(latent_dim,)),
            tf.keras.layers.Dense(units=7*7*32, activation=tf.nn.relu),
            tf.keras.layers.Reshape(target_shape=(7, 7, 32)),
            tf.keras.layers.Conv2DTranspose(
                filters=64, kernel_size=3, strides=2, padding='same',
                activation='relu'),
            tf.keras.layers.Conv2DTranspose(
                filters=32, kernel_size=3, strides=2, padding='same',
                activation='relu'),
            # No activation
            tf.keras.layers.Conv2DTranspose(
                filters=1, kernel_size=3, strides=1, padding='same'),
        ]
    )

  @tf.function
  def sample(self, eps=None):
    if eps is None:
      eps = tf.random.normal(shape=(100, self.latent_dim))
    return self.decode(eps, apply_sigmoid=True)

  def encode(self, x):
    mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)
    return mean, logvar

  def reparameterize(self, mean, logvar):
    eps = tf.random.normal(shape=mean.shape)
    return eps * tf.exp(logvar * .5) + mean

  def decode(self, z, apply_sigmoid=False):
    logits = self.decoder(z)
    if apply_sigmoid:
      probs = tf.sigmoid(logits)
      return probs
    return logits

損失関数とオプティマイザを定義します

VAEは、限界対数尤度の証拠下限(ELBO)を最大化することによってトレーニングします。

$$\log p(x) \ge \text{ELBO} = \mathbb{E}_{q(z|x)}\left[\log \frac{p(x, z)}{q(z|x)}\right].$$

実際には、この期待値の単一サンプルのモンテカルロ推定を最適化します。

$$\log p(x| z) + \log p(z) - \log q(z|x),$$

ここで、$ z $は$ q(z | x)$からサンプリングされます。

optimizer = tf.keras.optimizers.Adam(1e-4)


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.encode(x)
  z = model.reparameterize(mean, logvar)
  x_logit = model.decode(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.function
def train_step(model, x, optimizer):
  """Executes one training step and returns the loss.

  This function computes the loss and gradients, and uses the latter to
  update the model's parameters.
  """
  with tf.GradientTape() as tape:
    loss = compute_loss(model, x)
  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

トレーニング

  • データセットを反復処理することから始めます
  • 各反復中に、画像をエンコーダーに渡して、近似後部$ q(z | x)$の平均および対数分散パラメーターのセットを取得します。
  • 次に、再パラメーター化のトリックを適用して、$ q(z | x)$からサンプリングします。
  • 最後に、再パラメーター化されたサンプルをデコーダーに渡して、生成分布$ p(x | z)$のロジットを取得します。
  • 注:トレーニングセットに60kデータポイント、テストセットに10kデータポイントのkerasによって読み込まれたデータセットを使用するため、テストセットで得られたELBOは、LarochelleのMNISTの動的2値化を使用する文献で報告された結果よりもわずかに高くなります。

画像の生成

  • トレーニング後、いくつかの画像を生成する時が来ました
  • 単位ガウス事前分布$ p(z)$から潜在ベクトルのセットをサンプリングすることから始めます。
  • 次に、ジェネレータは潜在サンプル$ z $を観測のロジットに変換し、分布$ p(x | z)$を与えます。
  • ここでは、ベルヌーイ分布の確率をプロットします
epochs = 10
# set the dimensionality of the latent space to a plane for visualization later
latent_dim = 2
num_examples_to_generate = 16

# keeping the random vector constant for generation (prediction) so
# it will be easier to see the improvement.
random_vector_for_generation = tf.random.normal(
    shape=[num_examples_to_generate, latent_dim])
model = CVAE(latent_dim)
def generate_and_save_images(model, epoch, test_sample):
  mean, logvar = model.encode(test_sample)
  z = model.reparameterize(mean, logvar)
  predictions = model.sample(z)
  fig = plt.figure(figsize=(4, 4))

  for i in range(predictions.shape[0]):
    plt.subplot(4, 4, i + 1)
    plt.imshow(predictions[i, :, :, 0], cmap='gray')
    plt.axis('off')

  # tight_layout minimizes the overlap between 2 sub-plots
  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
  plt.show()
# Pick a sample of the test set for generating output images
assert batch_size >= num_examples_to_generate
for test_batch in test_dataset.take(1):
  test_sample = test_batch[0:num_examples_to_generate, :, :, :]
generate_and_save_images(model, 0, test_sample)

for epoch in range(1, epochs + 1):
  start_time = time.time()
  for train_x in train_dataset:
    train_step(model, train_x, optimizer)
  end_time = time.time()

  loss = tf.keras.metrics.Mean()
  for test_x in test_dataset:
    loss(compute_loss(model, test_x))
  elbo = -loss.result()
  display.clear_output(wait=False)
  print('Epoch: {}, Test set ELBO: {}, time elapse for current epoch: {}'
        .format(epoch, elbo, end_time - start_time))
  generate_and_save_images(model, epoch, test_sample)
Epoch: 10, Test set ELBO: -156.7376708984375, time elapse for current epoch: 4.276784658432007

png

最後のトレーニングエポックから生成された画像を表示します

def display_image(epoch_no):
  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))
plt.imshow(display_image(epoch))
plt.axis('off')  # Display images
(-0.5, 287.5, 287.5, -0.5)

png

保存されたすべての画像のアニメーションGIFを表示します

anim_file = 'cvae.gif'

with imageio.get_writer(anim_file, mode='I') as writer:
  filenames = glob.glob('image*.png')
  filenames = sorted(filenames)
  for filename in filenames:
    image = imageio.imread(filename)
    writer.append_data(image)
  image = imageio.imread(filename)
  writer.append_data(image)
import tensorflow_docs.vis.embed as embed
embed.embed_file(anim_file)

gif

潜在空間からの数字の2D多様体を表示します

以下のコードを実行すると、2D潜在空間全体で各桁が別の桁にモーフィングされた、さまざまな桁クラスの連続分布が表示されます。 TensorFlow Probabilityを使用して、潜在空間の標準正規分布を生成します。

def plot_latent_images(model, n, digit_size=28):
  """Plots n x n digit images decoded from the latent space."""

  norm = tfp.distributions.Normal(0, 1)
  grid_x = norm.quantile(np.linspace(0.05, 0.95, n))
  grid_y = norm.quantile(np.linspace(0.05, 0.95, n))
  image_width = digit_size*n
  image_height = image_width
  image = np.zeros((image_height, image_width))

  for i, yi in enumerate(grid_x):
    for j, xi in enumerate(grid_y):
      z = np.array([[xi, yi]])
      x_decoded = model.sample(z)
      digit = tf.reshape(x_decoded[0], (digit_size, digit_size))
      image[i * digit_size: (i + 1) * digit_size,
            j * digit_size: (j + 1) * digit_size] = digit.numpy()

  plt.figure(figsize=(10, 10))
  plt.imshow(image, cmap='Greys_r')
  plt.axis('Off')
  plt.show()
plot_latent_images(model, 20)

png

次のステップ

このチュートリアルでは、TensorFlowを使用して畳み込み変分オートエンコーダーを実装する方法を示しました。

次のステップとして、ネットワークサイズを増やすことにより、モデル出力の改善を試みることができます。たとえば、あなたが設定してみてください可能性がfilterのそれぞれのパラメータConv2DConv2DTranspose最終2D潜像プロットを生成するために、あなたが保持する必要があるだろうと512注意にレイヤーをlatent_dimまた2に、訓練時間が増加するであろうネットワークサイズが大きくなるにつれて。

CIFAR-10などの別のデータセットを使用してVAEを実装してみることもできます。

VAEは、さまざまなスタイルでさまざまな複雑さで実装できます。次のソースで追加の実装を見つけることができます。

VAEの詳細について詳しく知りたい場合は、 「変分オートエンコーダの概要」を参照してください。