#!/usr/bin/env python

# -*- coding: utf-8 -*-

"""

file: mnist_acgan.py

author: Luke de Oliveira (lukedeo@vaitech.io)

Train an Auxiliary Classifier Generative Adversarial Network (ACGAN) on the

MNIST dataset. See https://arxiv.org/abs/1610.09585 for more details.

You should start to see reasonable images after ~5 epochs, and good images

by ~15 epochs. You should use a GPU, as the convolution-heavy operations are

very slow on the CPU. Prefer the TensorFlow backend if you plan on iterating, as

the

Timings:

Hardware | Backend | Time / Epoch

-------------------------------------------

CPU | TF | 3 hrs

Titan X (maxwell) | TF | 4 min

Titan X (maxwell) | TH | 7 min

"""

from __future__ import print_function

from collections import defaultdict

import cPickle as pickle

from PIL import Image

from six.moves import range

from keras.backend import set_image_dim_ordering

from keras.datasets import mnist

from keras.layers import Input, Dense, Reshape, Flatten, Embedding, merge, Dropout

from keras.layers.advanced_activations import LeakyReLU

from keras.layers.convolutional import UpSampling2D, Convolution2D

from keras.models import Sequential, Model

from keras.optimizers import Adam

from keras.utils.generic_utils import Progbar

import numpy as np

np.random.seed(1337)

set_image_dim_ordering('th')

# nice, tabular formatting for progress messages

HEADING = '{0:<22s} | {1:4s} | {2:15s} | {3:5s}'

ROW_FMT = '{0:<22s} | {1:<4.2f} | {2:<15.2f} | {3:<5.2f}'

def adam_config():

# parameters suggested in https://arxiv.org/abs/1511.06434, slows down

# Adam for better convergence

return Adam(lr=0.0002, beta_1=0.5, beta_2=0.999, epsilon=1e-08)

def build_generator(latent_size):

return Model(input=[latent, image_class], output=fake_image)

def build_discriminator():

return Model(input=image, output=[fake, aux])

if __name__ == '__main__':

# batch and latent size taken from the paper

nb_epochs = 50

batch_size = 100

latent_size = 100

# build the discriminator

discriminator = build_discriminator()

discriminator.compile(adam_config(), ['binary_crossentropy',

'sparse_categorical_crossentropy'])

# build the generator

generator = build_generator(latent_size)

generator.compile(adam_config(), 'binary_crossentropy')

latent = Input(shape=(latent_size, ))

image_class = Input(shape=(1, 1), dtype='int32')

# get a fake image

fake = generator([latent, image_class])

# we only want to be able to train generation for the combined model

discriminator.trainable = False

fake, aux = discriminator(fake)

combined = Model(input=[latent, image_class], output=[fake, aux])

combined.compile(adam_config(), ['binary_crossentropy',

'sparse_categorical_crossentropy'])

discriminator.trainable = True

# get our mnist data, and force it to be of shape (..., 1, 28, 28) with

# range [-1, 1]

(X_train, y_train), (X_test, y_test) = mnist.load_data()

X_train = (X_train.astype(np.float32) - 127.5) / 127.5

X_train = np.expand_dims(X_train, axis=1)

X_test = (X_test.astype(np.float32) - 127.5) / 127.5

X_test = np.expand_dims(X_test, axis=1)

nb_train, nb_test = X_train.shape[0], X_test.shape[0]

train_history = defaultdict(list)

test_history = defaultdict(list)

for epoch in range(nb_epochs):

# print 'Epoch {} of {}'.format(epoch + 1, nb_epochs)

nb_batches = int(X_train.shape[0] / batch_size)

progress_bar = Progbar(target=nb_batches)

epoch_gen_loss = []

epoch_disc_loss = []

for index in range(nb_batches):

progress_bar.update(index)

# generate a new batch of noise

noise = np.random.uniform(-1, 1, (batch_size, latent_size))

# get a batch of real images

image_batch = X_train[index * batch_size:(index + 1) * batch_size]

label_batch = y_train[index * batch_size:(index + 1) * batch_size]

# sample some labels from p_c

sampled_labels = np.random.randint(0, 10, batch_size)

# generate a batch of fake images, using the generated labels as a

# conditioner. We reshape the sampled labels to be

# (batch_size, 1, 1) so that we can feed them into the embedding

# layer as a length one sequence

generated_images = generator.predict(

[noise, sampled_labels.reshape((-1, 1, 1))], verbose=0)

X = np.concatenate((image_batch, generated_images))

y = np.array([1] * batch_size + [0] * batch_size)

aux_y = np.concatenate((label_batch, sampled_labels), axis=0)

# see if the discriminator can figure itself out...

epoch_disc_loss.append(discriminator.train_on_batch(X, [y, aux_y]))

# make new noise. we generate 2 * batch size here such that we have

# the generator optimize over an identical number of images as the

# discriminator

noise = np.random.uniform(-1, 1, (2 * batch_size, latent_size))

sampled_labels = np.random.randint(0, 10, 2 * batch_size)

# we want to fix the discriminator and let the generator train to

# trick it

discriminator.trainable = False

# For the generator, we want all the {fake, not-fake} labels to say

# not-fake

trick = np.ones(2 * batch_size)

epoch_gen_loss.append(combined.train_on_batch(

[noise, sampled_labels.reshape((-1, 1, 1))], [trick, sampled_labels]))

discriminator.trainable = True

print('\nTesting for epoch {}:'.format(epoch + 1))

# evaluate the testing loss here

# generate a new batch of noise

noise = np.random.uniform(-1, 1, (nb_test, latent_size))

# sample some labels from p_c and generate images from them

sampled_labels = np.random.randint(0, 10, nb_test)

generated_images = generator.predict(

[noise, sampled_labels.reshape((-1, 1, 1))], verbose=False)

X = np.concatenate((X_test, generated_images))

y = np.array([1] * nb_test + [0] * nb_test)

aux_y = np.concatenate((y_test, sampled_labels), axis=0)

# see if the discriminator can figure itself out...

discriminator_test_loss = discriminator.evaluate(

X, [y, aux_y], verbose=False)

discriminator_train_loss = np.mean(np.array(epoch_disc_loss), axis=0)

# make new noise

noise = np.random.uniform(-1, 1, (2 * nb_test, latent_size))

sampled_labels = np.random.randint(0, 10, 2 * nb_test)

trick = np.ones(2 * nb_test)

generator_test_loss = combined.evaluate(

[noise, sampled_labels.reshape((-1, 1, 1))],

[trick, sampled_labels], verbose=False)

generator_train_loss = np.mean(np.array(epoch_gen_loss), axis=0)

# generate an epoch report on performance

train_history['generator'].append(generator_train_loss)

train_history['discriminator'].append(discriminator_train_loss)

test_history['generator'].append(generator_test_loss)

test_history['discriminator'].append(discriminator_test_loss)

print(HEADING.format('component', *discriminator.metrics_names))

print('-' * 65)

print(ROW_FMT.format('generator (train)',

*train_history['generator'][-1]))

print(ROW_FMT.format('generator (test)',

*test_history['generator'][-1]))

print(ROW_FMT.format('discriminator (train)',

*train_history['discriminator'][-1]))

print(ROW_FMT.format('discriminator (test)',

*test_history['discriminator'][-1]))

# save weights every epoch

generator.save_weights(

'params_generator_epoch_{0:03d}.hdf5'.format(epoch), True)

discriminator.save_weights(

'params_discriminator_epoch_{0:03d}.hdf5'.format(epoch), True)

# generate some digits to display

noise = np.random.uniform(-1, 1, (100, latent_size))

sampled_labels = np.array([

[i] * 10 for i in range(10)

]).reshape(-1, 1, 1)

# get a batch to display

generated_images = generator.predict(

[noise, sampled_labels], verbose=0)

# arrange them into a grid

img = (np.concatenate([r.reshape(-1, 28)

for r in np.split(generated_images, 10)

], axis=-1) * 127.5 + 127.5).astype(np.uint8)

Image.fromarray(img).save(

'plot_epoch_{0:03d}_generated.png'.format(epoch))

pickle.dump({'train': train_history, 'test': test_history},

open('acgan-history.pkl', 'wb'))