cvae.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Autoenconder variacional convolucional

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Este notebook demonstra como treinar um autoencoder variacional (VAE, na sigla em inglês) (1, 2) no dataset MNIST. Um VAE é uma versão probabilística do autoencoder, um modelo que recebe dados de entrada de dimensão alta e comprime-os em uma representação menor. Diferentemente de um autoencoder tradicional, que mapeia a entrada em um vetor latente, um VAE mapeia os dados de entrada nos parâmetros de uma distribuição de probabilidades, como a média e a variância gaussianas. Essa estratégia gera um espaço latente estruturado e contínuo, muito útil para a geração de imagens.

Configuração

!pip install tensorflow-probability

# to generate gifs
!pip install imageio
!pip install git+https://github.com/tensorflow/docs

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

Carregar o dataset MNIST

Cada imagem do MNIST é originalmente um vetor com 784 inteiros, sendo que cada um está entre 0 e 255 e representa a intensidade de um pixel. Modele cada pixel com uma distribuição de Bernoulli em nosso modelo e binarize estatisticamente o dataset.

(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

Usar o tf.data para dividir os dados em lotes e misturá-los

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))

Definir as redes do encoder e decoder com tf.keras.Sequential

Neste exemplo de VAE, use duas ConvNets pequenas para as redes do encoder e do decoder. Na literatura, essas redes também são chamadas de modelos de inferência/reconhecimento e generativo, respectivamente. Use o tf.keras.Sequential para simplificar a implementação. $x$ e $z$ denotam a observação e a variável latente, respectivamente, nas seguintes descrições.

Rede do encoder

Define a distribuição posterior aproximada $q(z|x)$, que recebe como entrada uma observação e gera como saída um conjunto de parâmetros para especificar a distribuição condicional da representação latente $z$. Neste exemplo, modele a distribuição como uma gaussiana diagonal, e a rede gera como saída a média e os parâmetros da variância logarítmica de uma distribuição gaussiana fatorial. Gere como a saída a variância logarítmica em vez da variância diretamente para proporcionar estabilidade numérica.

Rede do decoder

Define a distribuição condicional da observação $p(x|z)$, que recebe uma amostra latente $z$ como entrada e gera como saída os parâmetros de uma distribuição condicional da observação. Modele o prior $p(z)$ da distribuição latente como uma gaussiana unitária.

Truque de reparametrização

Para gerar uma amostra $z$ para o decoder durante o treinamento, você pode fazer uma amostragem da distribuição latente definida pelos parâmetros gerados como saída pelo encoder, dada uma observação de entrada $x$. Entretanto, essa operação de amostragem cria um gargalo, pois a retropropagação não pode fluir por um nó aleatório.

Para tratar esse problema, use um truque de reparametrização. Em nosso exemplo, você aproxima $z$ usando os parâmetros do decoder e outro parâmetro $\epsilon$ da seguinte forma:

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

em que $\mu$ e $\sigma$ representam a média e o desvio padrão de uma distribuição gaussiana, respectivamente. Eles podem ser derivados da saída do decoder. Podemos pensar no $\epsilon$ como um ruído aleatório usado para manter a estocasticidade de $z$. Gere $\epsilon$ a partir de uma distribuição normal padrão.

Agora, a variável latente $z$ é gerada por uma função de $\mu$, $\sigma$ e $\epsilon$, o que permite ao modelo retropropagar os gradientes no encoder por meio de $\mu$ e $\sigma$, respectivamente, ao mesmo tempo em que mantém a estocasticidade por meio de $\epsilon$.

Arquitetura da rede

Para a rede do encoder, use duas camadas convolucionais seguidas por uma camada totalmente conectada. Na rede do decoder, espelhe essa arquitetura usando uma camada totalmente conectada seguida por três camadas transpostas convolucionais (também conhecidas como camadas deconvolucionais em alguns contextos). É uma prática comum evitar o uso da normalização de lotes ao treinar VAEs, já que a estocasticidade adicional devido ao uso de minilotes pode agravar a instabilidade da estocasticidade oriunda da amostragem.

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

Definir a função de perda e o otimizador

Os VAEs são treinados pela maximização do limite inferior de evidência (ELBO, na sigla em inglês) na verossimilhança logarítmica marginal:

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

Na prática, otimize a estimativa de Monte Carlo de amostra única desta expectativa:

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

Observação: você também poderia computar analiticamente o termo KL, mas aqui você incorpora todos os três termos do estimador de Monte Carlo por questões de simplicidade.

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))

Treinamento

• Comece fazendo a iteração do dataset

• Durante cada iteração, passe a imagem para o encoder a fim de obter um conjunto de média e parâmetros de variância logarítmica do $q(z|x)$ posterior aproximado

• Depois, aplique o truque de reparametrização para fazer a amostragem a partir de $q(z|x)$

• Por fim, passe as amostras reparametrizadas para o decoder a fim de obter os logits da distribuição generativa $p(x|z)$

• Observação: como você utiliza o dataset carregado pelo Keras com 60 mil pontos de dados no conjunto de treinamento e 10 mil pontos de dados no conjunto de teste, o ELBO resultante no conjunto de teste é ligeiramente superior aos resultados indicados na literatura, que usa a binarização dinâmica do MNIST de Larochelle.

Geração de imagens

• Após fazer o treinamento, está na hora de gerar algumas imagens

• Comece fazendo a amostragem de um conjunto de vetores latentes a partir da distribuição de prior gaussiana unitária $p(z)$

• O gerador vai converter a amostra latente $z$ em logits da observação, gerando uma distribuição $p(x|z)$

• Agora, plote as probabilidades das distribuições de Bernoulli

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)

Exibição de uma imagem gerada na última época de treinamento

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

Exibição de um GIF animado de todas as imagens salvas

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)

Exibição de uma variedade bidimensional de dígitos do espaço latente

Ao executar o código abaixo, será exibida uma distribuição contínua das diferentes classes de dígitos, em que cada dígito transforma-se em outro no espaço latente bidimensional. Use a TensorFlow Probability para gerar uma distribuição normal padrão para o espaço latente.

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)

Próximos passos

Este tutorial demonstrou como implementar um autoencoder variacional convolucional usando o TensorFlow.

Agora, você pode tentar melhorar a saída do modelo aumentando o tamanho da rede. Por exemplo, você pode tentar definir os parâmetros de filter para cada uma das camadas Conv2D e Conv2DTranspose como 512. Para gerar o gráfico de imagem latente bidimensional, você precisaria manter latent_dim como 2. Além disso, o tempo de treinamento aumenta à medida que o tamanho da rede aumenta.

Você também pode tentar implementar um VAE usando um dataset diferente, como o CIFAR-10.

Os VAEs podem ser implementados com diferentes estilos e complexidades variadas. Confira outras implementações nas seguintes fontes:

• Autoencoder variacional (keras.io)

• Exemplo de VAE do guia "Como escrever camadas e modelos personalizados" (tensorflow.org)

• Camadas probabilísticas: autoencoder variacional

Se quiser saber mais sobre os detalhes de VAEs, confira Introdução aos autoencoders variacionais.