dcgan.ipynb

Kernel: Python 3

Copyright 2019 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.

Rede Adversária Generativa Convolucional Profunda

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Este tutorial demonstra como gerar imagens de dígitos escritos à mão usando uma Rede Adversária Generativa Convolucional Profunda (DCGAN, na sigla em inglês). O código é escrito usando a API Sequential do Keras, com um loop de treinamento tf.GradientTape.

O que são GANs?

As Redes Adversárias Generativas (GANs, na sigla em inglês) são uma das ideias mais interessantes da ciência da computação da atualidade. Dois modelos são treinados simultaneamente por um processo adversário. Um gerador (“o artista”) aprende a criar imagens que parecem reais, enquanto um discriminador (“o crítico de arte”) aprende a diferenciar imagens reais de falsas.

Durante o treinamento, o gerador fica cada vez melhor em criar imagens que parecem reais, enquanto o discriminador fica cada vez melhor em identificá-las. O processo alcança o equilíbrio quando o discriminador não consegue mais distinguir imagens reais de falsas.

Este notebook demonstra esse processo no dataset MNIST. O GIF animado abaixo mostra uma série de imagens produzidas pelo gerador após ser treinado com 50 épocas. As imagens começam como um ruído aleatório e se parecem cada vez mais com dígitos escritos à mão ao longo do tempo.

Para saber mais sobre GANs, confira o curso Introdução ao aprendizado profundo do MIT.

Configuração

import tensorflow as tf

tf.__version__

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

import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
from tensorflow.keras import layers
import time

from IPython import display

Carregar e preparar o dataset

Você usará o dataset MNIST para treinar o gerador e o discriminador. O gerador vai gerar dígitos escritos à mão que parecem os dados do MNIST.

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

train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
train_images = (train_images - 127.5) / 127.5  # Normalize the images to [-1, 1]

BUFFER_SIZE = 60000
BATCH_SIZE = 256

# Batch and shuffle the data
train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

Criar os modelos

Tanto o gerador quanto o discriminador são definidos usando a API Sequential do Keras.

O gerador

O gerador usa camadas do tf.keras.layers.Conv2DTranspose (upsampling) para produzir uma imagem a partir de uma semente (ruído aleatório). Comece com uma camada Dense que recebe essa semente como entrada, depois faça o upsampling diversas vezes até conseguir o tamanho de imagem desejado de 28x28x1. Observe a ativação do tf.keras.layers.LeakyReLU para cada camada, exceto pela camada de saída, que usa tanh.

def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 28, 28, 1)

    return model

Use o gerador (ainda não treinado) para criar uma imagem.

generator = make_generator_model()

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

plt.imshow(generated_image[0, :, :, 0], cmap='gray')

O discriminador

O discriminador é um classificador de imagens baseado em uma CNN.

def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                                     input_shape=[28, 28, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(1))

    return model

Use o discriminador (ainda não treinado) para classificar as imagens geradas como reais ou falsas. O modelo será treinado para gerar como saída valores positivos para imagens reais e valores negativos para imagens falsas.

discriminator = make_discriminator_model()
decision = discriminator(generated_image)
print (decision)

Definir a perda e os otimizadores

Defina funções de perda e otimizadores para os dois modelos.

# This method returns a helper function to compute cross entropy loss
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)

Perda do discriminador

Este método quantifica em que nível o discriminador consegue distinguir imagens reais de falsas. Ele compara as previsões de imagens reais do discriminador com um array de 1s e as previsões de imagens falsas (geradas) do discriminador com um array de 0s.

def discriminator_loss(real_output, fake_output):
    real_loss = cross_entropy(tf.ones_like(real_output), real_output)
    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    return total_loss

Perda do gerador

A perda do gerador quantifica em que nível ele consegue enganar o discriminador. Intuitivamente, se o gerador estiver com bom desempenho, o discriminador classificará as imagens falsas como reais (ou 1). Agora, compare as decisões do discriminador sobre as imagens geradas com um array de 1s.

def generator_loss(fake_output):
    return cross_entropy(tf.ones_like(fake_output), fake_output)

Os otimizadores do discriminador e do gerador são diferentes, já que você treinará duas redes separadamente.

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

Salvar checkpoints

Este notebook também demonstra como salvar e restaurar modelos, o que pode ser útil caso uma tarefa de treinamento demorada seja interrompida.

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

Definir o loop de treinamento

EPOCHS = 50
noise_dim = 100
num_examples_to_generate = 16

# You will reuse this seed overtime (so it's easier)
# to visualize progress in the animated GIF)
seed = tf.random.normal([num_examples_to_generate, noise_dim])

O loop de treinamento começa com o gerador recebendo uma semente aleatória como entrada. Essa semente é usada para gerar uma imagem. Em seguida, o discriminador é usado para classificar imagens reais (obtidas no conjunto de treinamento) e imagens falsas (geradas pelo gerador). A perda é calculada para cada modelo, e os gradientes são usados para atualizar o gerador e o discriminador.

# Notice the use of `tf.function`
# This annotation causes the function to be "compiled".
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      generated_images = generator(noise, training=True)

      real_output = discriminator(images, training=True)
      fake_output = discriminator(generated_images, training=True)

      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

def train(dataset, epochs):
  for epoch in range(epochs):
    start = time.time()

    for image_batch in dataset:
      train_step(image_batch)

    # Produce images for the GIF as you go
    display.clear_output(wait=True)
    generate_and_save_images(generator,
                             epoch + 1,
                             seed)

    # Save the model every 15 epochs
    if (epoch + 1) % 15 == 0:
      checkpoint.save(file_prefix = checkpoint_prefix)

    print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))

  # Generate after the final epoch
  display.clear_output(wait=True)
  generate_and_save_images(generator,
                           epochs,
                           seed)

Gerar e salvar imagens

def generate_and_save_images(model, epoch, test_input):
  # Notice `training` is set to False.
  # This is so all layers run in inference mode (batchnorm).
  predictions = model(test_input, training=False)

  fig = plt.figure(figsize=(4, 4))

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

  plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
  plt.show()

Treinar o modelo

Faça uma chamada ao método train() definido acima para treinar o gerador e o discriminador simultaneamente. Treinar GANs pode ser complicado. É importante que o gerador e o discriminador não sobrecarreguem um ao outro (por exemplo, não devem fazer o treinamento a uma taxa similar).

No começo do treinamento, as imagens geradas parecem ruído aleatório. À medida que o treinamento progride, os dígitos gerados parecerão cada vez mais reais. Após cerca de 50 épocas, eles parecem dígitos do MNIST. Isso pode levar um minuto por época com as configurações padrão do Colab.

train(train_dataset, EPOCHS)

Restaure o último checkpoint.

checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

Criar um GIF

# Display a single image using the epoch number
def display_image(epoch_no):
  return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))

display_image(EPOCHS)

Use o imageio para criar um GIF animado usando as imagens salvas durante o treinamento.

anim_file = 'dcgan.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)

Próximos passos

Este tutorial demonstrou o código necessário para escrever e treinar uma GAN. Agora, talvez você queira experimentar um dataset diferente, por exemplo, o dataset Atributos de rostos de celebridades em grande escala (CelebA), disponível no Kaggle. Para saber mais sobre GANs, confira o Tutorial do NIPS 2016: Redes Adversárias Generativas.