GitHub Repository: tensorflow/docs-l10n
Path: blob/master/site/es-419/tutorials/generative/cvae.ipynb
²⁵¹¹⁸ views

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

In [ ]:

#@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.

Autocodificador variacional convolucional

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar el bloc de notas

En este cuaderno se enseña cómo entrenar un autocodificador variacional (VAE por sus siglas en inglés) (1, 2) en el conjunto de datos MNIST. Un VAE es un enfoque probabilístico del autocodificador, es un modelo que toma datos de entrada de alta dimensión y los comprime en una representación más pequeña. A diferencia de un autocodificador tradicional, que mapea la entrada en un vector latente, un VAE mapea los datos de entrada en parámetros de distribución de probabilidad, tales como la media y la varianza de una gaussiana. Este enfoque produce un espacio de latente continuo y estructurado, que sirve para generar una imagen.

Preparación

In [ ]:

!pip install tensorflow-probability

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

In [ ]:

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

Cargar el conjunto de datos MNIST

Cada imagen de MNIST es originalmente un vector de 784 enteros, cada uno de entre 0 y 255 y representan la intensidad de un pixel. Modele cada pixel con la distribución de Bernoulli en nuestro modelo y convierta el conjunto de datos en binario de forma estática.

In [ ]:

(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()

In [ ]:

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)

In [ ]:

train_size = 60000
batch_size = 32
test_size = 10000

Usar tf.data para poner los datos en lotes y en orden aleatorio

In [ ]:

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 las redes del codificador y del decodificador con tf.keras.Sequential

En este ejemplo de VAE, use dos ConvNets pequeños para las redes del codificador y del decodificador. En la documentación, también se las conoce como inferencia/reconocimiento y modelos generativos respectivamente. Use tf.keras.Sequential para simplificar la implementación. Dejemos que $x$ y $z$ denoten la observación y la variable de latente respectivamente en las siguientes descripciones.

Red del codificador

Define la distribución posterior aproximada de $q(z|x)$ , que toma como entrada una observación y tiene como salida un conjunto de parámetros para especificar la distribución condicional de la representación latente $z$ . En este ejemplo, simplemente modele la distribución como una gaussiana diagonal y la salida de la red será los parámetros de la media y de la varianza logarítmica de una gaussiana factorizada. Tiene como salida directamente el logaritmo de la varianza logarítmica en vez de la varianza para conservar la estabilidad numérica.

Red del decodificador

Define la distribución condicional de la observación $p(x|z)$ , que toma una muestra de latente $z$ como entrada y tiene como salida los parámetros para una distribución condicional de la observación. Modele la distribución latente antes de $p(z)$ como una gaussiana de unidad.

Engaño de reparametrización

Para generar una muestra $z$ para el decodificador durante el entrenamiento, puede tomar una muestra de la distribución latente definida por los parámetros que salieron del codificador, al recibir una observación de entrada $x$ . Sin embargo, esta operación de muestra crea un cuello de botella porque la retropropagación no puede fluir a través de un nodo aleatorio.

Para arreglar esto, use un engaño de reparametrización. En nuestro ejemplo, usted aproximará $z$ mediante el uso de los parámetros del decodificador y otro parámetro $\epsilon$ de la siguiente manera:

z = \mu + \sigma \odot \epsilon

donde $\mu$ y $\sigma$ representan la desviación media y estándar de una distribución gaussiana respectivamente. Pueden derivarse de la salida del decodificador. El $\epsilon$ puede ser considerado como un ruido aleatorio para conservar la estocasticidad de $z$ . Genere $\epsilon$ a partir de una distribución normal estándar.

La variable latente $z$ ahora se genera mediante la función de $\mu$ , $\sigma$ y $\epsilon$ , que permitirá que el modelo retropropague gradientes en el codificador a través de $\mu$ y $\sigma$ respectivamente, mientras conserva la estocasticidad a través de $\epsilon$ .

La arquitectura de la red

Para la red del codificador, use dos capas convolucionales seguidas de una capa completamente conectada. En la red del decodificador, replique esta arquitectura usando una capa completamente conectada seguida de tres capas de convolución transpuesta (también conocidas como capas deconvolucionales en algunos contextos). Nota, es común evitar el uso de normalización de lotes cuando se entrenan las VAE, ya que la estocasticidad adicional causada por el uso de mini lotes puede agravar la inestabilidad además de la estocasticidad de las muestras.

In [ ]:

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 la función de pérdida y el optimizador

Las VAE entrenan al maximizar el límite inferior de evidencia (ELBO, por sus siglas en inglés) en la probabilidad logarítmica marginal:

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

En la práctica, optimice el cálculo estimado de Monte Carlo de la muestra única del siguiente valor esperado:

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

Nota: También puede calcular analíticamente el término KL, pero en este caso puede incorporar los tres términos en el estimador de Monte Carlo para que sea más simple.

In [ ]:

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

Entrenamiento

Primero, iteramos el conjunto de datos
Durante cada iteración, pase la imagen al codificador para obtener un conjunto de parámetros de la media y de la varianza de logaritmos de la posterior aproximada $q(z|x)$
luego aplique el engaño de reparametrización para tomar una muestra de $q(z|x)$
Por último, pase las muestras reparametrizadas al decodificador para obtener los logit de la distribución generativa $p(x|z)$
Nota: Ya que se usó el conjunto de datos cargado mediante Keras con 60k puntos de datos en el conjunto de entrenamiento y 10k puntos de datos en el conjunto de prueba, nuestro ELBO resultante en el conjunto de prueba es un poco más alto que los resultados registrados en la documentación que usa una binarización dinámica de la MNIST de Larochelle.

Generar imágenes

Después del entrenamiento, es hora de generar algunas imágenes
Para empezar, tome una muestra del conjunto de vectores latentes de la gaussiana de unidad antes de la distribución $p(z)$
Luego de esto, el generador convertirá la muestra latente $z$ en logits de observación, lo que resultará en una distribución $p(x|z)$
Aquí, puede trazar las probabilidades de las distribuciones de Bernoulli

In [ ]:

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)

In [ ]:

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

In [ ]:

# 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, :, :, :]

In [ ]:

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)

Mostrar una imagen generada en la última época del entrenamiento

In [ ]:

def display_image(epoch_no):
  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))

In [ ]:

plt.imshow(display_image(epoch))
plt.axis('off')  # Display images

Mostrar un GIF animado de todas las imágenes guardadas

In [ ]:

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)

In [ ]:

import tensorflow_docs.vis.embed as embed
embed.embed_file(anim_file)

Mostrar una variedad bidimensional de cifras del espacio latente

Al ejecutar el siguiente código, se mostrará una distribución continua de diferentes clases de cifras, con cada cifra transformándose en otra en todo el espacio latente bidimensional. Use TensorFlow Probability para generar una distribución estándar normal para el espacio latente.

In [ ]:

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

In [ ]:

plot_latent_images(model, 20)

Próximos pasos

En este tutorial, se explicó cómo implementar un autocodificador variacional convolucional con TensorFlow.

Su próximo paso puede ser intentar mejorar la salida del modelo al aumentar el tamaño de la red. Por ejemplo, puede intentar configurar los parámetros de filter para cada una de las capas Conv2D y Conv2DTranspose a 512. Tenga en cuenta que, para generar el trazado bidimensional de la imagen latente, deberá conservar latent_dim en 2. Además, el tiempo de entrenamiento se reducirá a medida que aumente el tamaño de la red.

También puede intentar implementar una VAE con un conjunto de datos diferente, como el CIFAR-10.

Las VAE pueden implementarse con diferentes estilos y diversas complejidades. Puede encontrar más implementaciones en los siguientes recursos:

Si le gustaría obtener más información sobre los detalles de las VAE, vea Una introducción a los autocodificadores variacionales.