Train an Auxiliary Classifier GAN (ACGAN) on the MNIST dataset.

More details on Auxiliary Classifier GANs.

You should start to see reasonable images after ~5 epochs, and good imagesby ~15 epochs. You should use a GPU, as the convolution-heavy operations arevery slow on the CPU. Prefer the TensorFlow backend if you plan on iterating,as the compilation time can be a blocker using Theano.

Timings:

HardwareBackendTime / Epoch
CPUTF3 hrs
Titan X (maxwell)TF4 min
Titan X (maxwell)TH7 min

Consult Auxiliary Classifier Generative Adversarial Networks in Keras for more information and example output.

  1. from __future__ import print_function
  3. from collections import defaultdict
  4. try:
  5.     import cPickle as pickle
  6. except ImportError:
  7.     import pickle
  8. from PIL import Image
  10. from six.moves import range
  12. from keras.datasets import mnist
  13. from keras import layers
  14. from keras.layers import Input, Dense, Reshape, Flatten, Embedding, Dropout
  15. from keras.layers import BatchNormalization
  16. from keras.layers.advanced_activations import LeakyReLU
  17. from keras.layers.convolutional import Conv2DTranspose, Conv2D
  18. from keras.models import Sequential, Model
  19. from keras.optimizers import Adam
  20. from keras.utils.generic_utils import Progbar
  21. import numpy as np
  23. np.random.seed(1337)
  24. num_classes = 10
  27. def build_generator(latent_size):
  28.     # we will map a pair of (z, L), where z is a latent vector and L is a
  29.     # label drawn from P_c, to image space (..., 28, 28, 1)
  30.     cnn = Sequential()
  32.     cnn.add(Dense(3 * 3 * 384, input_dim=latent_size, activation='relu'))
  33.     cnn.add(Reshape((3, 3, 384)))
  35.     # upsample to (7, 7, ...)
  36.     cnn.add(Conv2DTranspose(192, 5, strides=1, padding='valid',
  37.                             activation='relu',
  38.                             kernel_initializer='glorot_normal'))
  39.     cnn.add(BatchNormalization())
  41.     # upsample to (14, 14, ...)
  42.     cnn.add(Conv2DTranspose(96, 5, strides=2, padding='same',
  43.                             activation='relu',
  44.                             kernel_initializer='glorot_normal'))
  45.     cnn.add(BatchNormalization())
  47.     # upsample to (28, 28, ...)
  48.     cnn.add(Conv2DTranspose(1, 5, strides=2, padding='same',
  49.                             activation='tanh',
  50.                             kernel_initializer='glorot_normal'))
  52.     # this is the z space commonly referred to in GAN papers
  53.     latent = Input(shape=(latent_size, ))
  55.     # this will be our label
  56.     image_class = Input(shape=(1,), dtype='int32')
  58.     cls = Embedding(num_classes, latent_size,
  59.                     embeddings_initializer='glorot_normal')(image_class)
  61.     # hadamard product between z-space and a class conditional embedding
  62.     h = layers.multiply([latent, cls])
  64.     fake_image = cnn(h)
  66.     return Model([latent, image_class], fake_image)
  69. def build_discriminator():
  70.     # build a relatively standard conv net, with LeakyReLUs as suggested in
  71.     # the reference paper
  72.     cnn = Sequential()
  74.     cnn.add(Conv2D(32, 3, padding='same', strides=2,
  75.                    input_shape=(28, 28, 1)))
  76.     cnn.add(LeakyReLU(0.2))
  77.     cnn.add(Dropout(0.3))
  79.     cnn.add(Conv2D(64, 3, padding='same', strides=1))
  80.     cnn.add(LeakyReLU(0.2))
  81.     cnn.add(Dropout(0.3))
  83.     cnn.add(Conv2D(128, 3, padding='same', strides=2))
  84.     cnn.add(LeakyReLU(0.2))
  85.     cnn.add(Dropout(0.3))
  87.     cnn.add(Conv2D(256, 3, padding='same', strides=1))
  88.     cnn.add(LeakyReLU(0.2))
  89.     cnn.add(Dropout(0.3))
  91.     cnn.add(Flatten())
  93.     image = Input(shape=(28, 28, 1))
  95.     features = cnn(image)
  97.     # first output (name=generation) is whether or not the discriminator
  98.     # thinks the image that is being shown is fake, and the second output
  99.     # (name=auxiliary) is the class that the discriminator thinks the image
  100.     # belongs to.
  101.     fake = Dense(1, activation='sigmoid', name='generation')(features)
  102.     aux = Dense(num_classes, activation='softmax', name='auxiliary')(features)
  104.     return Model(image, [fake, aux])
  107. if __name__ == '__main__':
  108.     # batch and latent size taken from the paper
  109.     epochs = 100
  110.     batch_size = 100
  111.     latent_size = 100
  113.     # Adam parameters suggested in https://arxiv.org/abs/1511.06434
  114.     adam_lr = 0.0002
  115.     adam_beta_1 = 0.5
  117.     # build the discriminator
  118.     print('Discriminator model:')
  119.     discriminator = build_discriminator()
  120.     discriminator.compile(
  121.         optimizer=Adam(learning_rate=adam_lr, beta_1=adam_beta_1),
  122.         loss=['binary_crossentropy', 'sparse_categorical_crossentropy']
  123.     )
  124.     discriminator.summary()
  126.     # build the generator
  127.     generator = build_generator(latent_size)
  129.     latent = Input(shape=(latent_size, ))
  130.     image_class = Input(shape=(1,), dtype='int32')
  132.     # get a fake image
  133.     fake = generator([latent, image_class])
  135.     # we only want to be able to train generation for the combined model
  136.     discriminator.trainable = False
  137.     fake, aux = discriminator(fake)
  138.     combined = Model([latent, image_class], [fake, aux])
  140.     print('Combined model:')
  141.     combined.compile(
  142.         optimizer=Adam(learning_rate=adam_lr, beta_1=adam_beta_1),
  143.         loss=['binary_crossentropy', 'sparse_categorical_crossentropy']
  144.     )
  145.     combined.summary()
  147.     # get our mnist data, and force it to be of shape (..., 28, 28, 1) with
  148.     # range [-1, 1]
  149.     (x_train, y_train), (x_test, y_test) = mnist.load_data()
  150.     x_train = (x_train.astype(np.float32) - 127.5) / 127.5
  151.     x_train = np.expand_dims(x_train, axis=-1)
  153.     x_test = (x_test.astype(np.float32) - 127.5) / 127.5
  154.     x_test = np.expand_dims(x_test, axis=-1)
  156.     num_train, num_test = x_train.shape[0], x_test.shape[0]
  158.     train_history = defaultdict(list)
  159.     test_history = defaultdict(list)
  161.     for epoch in range(1, epochs + 1):
  162.         print('Epoch {}/{}'.format(epoch, epochs))
  164.         num_batches = int(np.ceil(x_train.shape[0] / float(batch_size)))
  165.         progress_bar = Progbar(target=num_batches)
  167.         epoch_gen_loss = []
  168.         epoch_disc_loss = []
  170.         for index in range(num_batches):
  171.             # get a batch of real images
  172.             image_batch = x_train[index * batch_size:(index + 1) * batch_size]
  173.             label_batch = y_train[index * batch_size:(index + 1) * batch_size]
  175.             # generate a new batch of noise
  176.             noise = np.random.uniform(-1, 1, (len(image_batch), latent_size))
  178.             # sample some labels from p_c
  179.             sampled_labels = np.random.randint(0, num_classes, len(image_batch))
  181.             # generate a batch of fake images, using the generated labels as a
  182.             # conditioner. We reshape the sampled labels to be
  183.             # (len(image_batch), 1) so that we can feed them into the embedding
  184.             # layer as a length one sequence
  185.             generated_images = generator.predict(
  186.                 [noise, sampled_labels.reshape((-1, 1))], verbose=0)
  188.             x = np.concatenate((image_batch, generated_images))
  190.             # use one-sided soft real/fake labels
  191.             # Salimans et al., 2016
  192.             # https://arxiv.org/pdf/1606.03498.pdf (Section 3.4)
  193.             soft_zero, soft_one = 0, 0.95
  194.             y = np.array(
  195.                 [soft_one] * len(image_batch) + [soft_zero] * len(image_batch))
  196.             aux_y = np.concatenate((label_batch, sampled_labels), axis=0)
  198.             # we don't want the discriminator to also maximize the classification
  199.             # accuracy of the auxiliary classifier on generated images, so we
  200.             # don't train discriminator to produce class labels for generated
  201.             # images (see https://openreview.net/forum?id=rJXTf9Bxg).
  202.             # To preserve sum of sample weights for the auxiliary classifier,
  203.             # we assign sample weight of 2 to the real images.
  204.             disc_sample_weight = [np.ones(2 * len(image_batch)),
  205.                                   np.concatenate((np.ones(len(image_batch)) * 2,
  206.                                                   np.zeros(len(image_batch))))]
  208.             # see if the discriminator can figure itself out...
  209.             epoch_disc_loss.append(discriminator.train_on_batch(
  210.                 x, [y, aux_y], sample_weight=disc_sample_weight))
  212.             # make new noise. we generate 2 * batch size here such that we have
  213.             # the generator optimize over an identical number of images as the
  214.             # discriminator
  215.             noise = np.random.uniform(-1, 1, (2 * len(image_batch), latent_size))
  216.             sampled_labels = np.random.randint(0, num_classes, 2 * len(image_batch))
  218.             # we want to train the generator to trick the discriminator
  219.             # For the generator, we want all the {fake, not-fake} labels to say
  220.             # not-fake
  221.             trick = np.ones(2 * len(image_batch)) * soft_one
  223.             epoch_gen_loss.append(combined.train_on_batch(
  224.                 [noise, sampled_labels.reshape((-1, 1))],
  225.                 [trick, sampled_labels]))
  227.             progress_bar.update(index + 1)
  229.         print('Testing for epoch {}:'.format(epoch))
  231.         # evaluate the testing loss here
  233.         # generate a new batch of noise
  234.         noise = np.random.uniform(-1, 1, (num_test, latent_size))
  236.         # sample some labels from p_c and generate images from them
  237.         sampled_labels = np.random.randint(0, num_classes, num_test)
  238.         generated_images = generator.predict(
  239.             [noise, sampled_labels.reshape((-1, 1))], verbose=False)
  241.         x = np.concatenate((x_test, generated_images))
  242.         y = np.array([1] * num_test + [0] * num_test)
  243.         aux_y = np.concatenate((y_test, sampled_labels), axis=0)
  245.         # see if the discriminator can figure itself out...
  246.         discriminator_test_loss = discriminator.evaluate(
  247.             x, [y, aux_y], verbose=False)
  249.         discriminator_train_loss = np.mean(np.array(epoch_disc_loss), axis=0)
  251.         # make new noise
  252.         noise = np.random.uniform(-1, 1, (2 * num_test, latent_size))
  253.         sampled_labels = np.random.randint(0, num_classes, 2 * num_test)
  255.         trick = np.ones(2 * num_test)
  257.         generator_test_loss = combined.evaluate(
  258.             [noise, sampled_labels.reshape((-1, 1))],
  259.             [trick, sampled_labels], verbose=False)
  261.         generator_train_loss = np.mean(np.array(epoch_gen_loss), axis=0)
  263.         # generate an epoch report on performance
  264.         train_history['generator'].append(generator_train_loss)
  265.         train_history['discriminator'].append(discriminator_train_loss)
  267.         test_history['generator'].append(generator_test_loss)
  268.         test_history['discriminator'].append(discriminator_test_loss)
  270.         print('{0:<22s} | {1:4s} | {2:15s} | {3:5s}'.format(
  271.             'component', *discriminator.metrics_names))
  272.         print('-' * 65)
  274.         ROW_FMT = '{0:<22s} | {1:<4.2f} | {2:<15.4f} | {3:<5.4f}'
  275.         print(ROW_FMT.format('generator (train)',
  276.                              *train_history['generator'][-1]))
  277.         print(ROW_FMT.format('generator (test)',
  278.                              *test_history['generator'][-1]))
  279.         print(ROW_FMT.format('discriminator (train)',
  280.                              *train_history['discriminator'][-1]))
  281.         print(ROW_FMT.format('discriminator (test)',
  282.                              *test_history['discriminator'][-1]))
  284.         # save weights every epoch
  285.         generator.save_weights(
  286.             'params_generator_epoch_{0:03d}.hdf5'.format(epoch), True)
  287.         discriminator.save_weights(
  288.             'params_discriminator_epoch_{0:03d}.hdf5'.format(epoch), True)
  290.         # generate some digits to display
  291.         num_rows = 40
  292.         noise = np.tile(np.random.uniform(-1, 1, (num_rows, latent_size)),
  293.                         (num_classes, 1))
  295.         sampled_labels = np.array([
  296.             [i] * num_rows for i in range(num_classes)
  297.         ]).reshape(-1, 1)
  299.         # get a batch to display
  300.         generated_images = generator.predict(
  301.             [noise, sampled_labels], verbose=0)
  303.         # prepare real images sorted by class label
  304.         real_labels = y_train[(epoch - 1) * num_rows * num_classes:
  305.                               epoch * num_rows * num_classes]
  306.         indices = np.argsort(real_labels, axis=0)
  307.         real_images = x_train[(epoch - 1) * num_rows * num_classes:
  308.                               epoch * num_rows * num_classes][indices]
  310.         # display generated images, white separator, real images
  311.         img = np.concatenate(
  312.             (generated_images,
  313.              np.repeat(np.ones_like(x_train[:1]), num_rows, axis=0),
  314.              real_images))
  316.         # arrange them into a grid
  317.         img = (np.concatenate([r.reshape(-1, 28)
  318.                                for r in np.split(img, 2 * num_classes + 1)
  319.                                ], axis=-1) * 127.5 + 127.5).astype(np.uint8)
  321.         Image.fromarray(img).save(
  322.             'plot_epoch_{0:03d}_generated.png'.format(epoch))
  324.     with open('acgan-history.pkl', 'wb') as f:
  325.         pickle.dump({'train': train_history, 'test': test_history}, f)