利用tensorflow训练简单的生成对抗网络GAN

import math
import os
import hashlib
from urllib.request import urlretrieve
import zipfile
import gzip
import shutil

data_dir = './data'

def download_extract(database_name, data_path):
     """
     Download and extract database
     :param database_name: Database name
     """
     DATASET_CELEBA_NAME = 'celeba'
     DATASET_MNIST_NAME = 'mnist'

     if database_name == DATASET_CELEBA_NAME:
         url = 'https://s3-us-west-1.amazonaws.com/udacity-dlnfd/datasets/celeba.zip'
         hash_code = '00d2c5bc6d35e252742224ab0c1e8fcb'
         extract_path = os.path.join(data_path, 'img_align_celeba')
         save_path = os.path.join(data_path, 'celeba.zip')
         extract_fn = _unzip
     elif database_name == DATASET_MNIST_NAME:
         url = 'http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz'
         hash_code = 'f68b3c2dcbeaaa9fbdd348bbdeb94873'
         extract_path = os.path.join(data_path, 'mnist')
         save_path = os.path.join(data_path, 'train-images-idx3-ubyte.gz')
         extract_fn = _ungzip

     if os.path.exists(extract_path):
         print('Found {} Data'.format(database_name))
         return

     if not os.path.exists(data_path):
         os.makedirs(data_path)

     if not os.path.exists(save_path):
         with DLProgress(unit='B', unit_scale=True, miniters=1, desc='Downloading {}'.format(database_name)) as pbar:
             urlretrieve(
                 url,
                 save_path,
                 pbar.hook)

     assert hashlib.md5(open(save_path, 'rb').read()).hexdigest() == hash_code, \
         '{} file is corrupted.  Remove the file and try again.'.format(save_path)

     os.makedirs(extract_path)
     try:
         extract_fn(save_path, extract_path, database_name, data_path)
     except Exception as err:
         shutil.rmtree(extract_path)  # Remove extraction folder if there is an error
         raise err

     # Remove compressed data
     os.remove(save_path)

# download mnist
download_extract('mnist', data_dir)
# download celeba
download_extract('celeba', data_dir

 # the number of images
  show_n_images =16
  
  %matplotlib inline
  import os
  from glob import glob
  from matplotlib import pyplot
  
  def get_batch(image_files, width, height, mode):
     data_batch = np.array(
         [get_image(sample_file, width, height, mode) for sample_file in image_files]).astype(np.float32)
 
     # Make sure the images are in 4 dimensions
     if len(data_batch.shape) < 4:
         data_batch = data_batch.reshape(data_batch.shape + (1,))
 
     return data_batch
 
 def images_square_grid(images, mode):
     # Get maximum size for square grid of images
     save_size = math.floor(np.sqrt(images.shape[0]))
 
     # Scale to 0-255
     images = (((images - images.min()) * 255) / (images.max() - images.min())).astype(np.uint8)
 
     # Put images in a square arrangement
     images_in_square = np.reshape(
             images[:save_size*save_size],
             (save_size, save_size, images.shape[1], images.shape[2], images.shape[3]))
     if mode == 'L':
         images_in_square = np.squeeze(images_in_square, 4)
 
     # Combine images to grid image
     new_im = Image.new(mode, (images.shape[1] * save_size, images.shape[2] * save_size))
     for col_i, col_images in enumerate(images_in_square):
         for image_i, image in enumerate(col_images):
             im = Image.fromarray(image, mode)
             new_im.paste(im, (col_i * images.shape[1], image_i * images.shape[2]))
 
     return new_im
 
 mnist_images = get_batch(glob(os.path.join(data_dir, 'mnist/*.jpg'))[:show_n_images], 28, 28, 'L')
 pyplot.imshow(images_square_grid(mnist_images, 'L'), cmap='gray')

show_n_images = 9
 
mnist_images = get_batch(glob(os.path.join(data_dir, 'img_align_celeba/*.jpg'))[:show_n_images], 28, 28, 'RGB')
pyplot.imshow(images_square_grid(mnist_images, 'RGB'))

from distutils.version import LooseVersion
import warnings
import tensorflow as tf

# Check TensorFlow Version
assert LooseVersion(tf.__version__) >= LooseVersion('1.0'), 'Please use TensorFlow version 1.0 or newer.  You are using {}'.format(tf.__version__)
print('TensorFlow Version: {}'.format(tf.__version__))

# Check for a GPU
if not tf.test.gpu_device_name():
    warnings.warn('No GPU found. Please use a GPU to train your neural network.')
else:
    print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))

def model_inputs(image_width, image_height, image_channels, z_dim):
    ## Real imag
    inputs_real = tf.placeholder(tf.float32,(None, image_width,image_height,image_channels), name = 'input_real')

    ## input z

    inputs_z = tf.placeholder(tf.float32,(None, z_dim), name='input_z')

    ## Learning rate
    learning_rate = tf.placeholder(tf.float32, name = 'lr')

    return inputs_real, inputs_z, learning_rate

def discriminator(images, reuse=False):
    """
    Create the discriminator network
    :param images: Tensor of input image(s)
    :param reuse: Boolean if the weights should be reused
    :return: Tuple of (tensor output of the discriminator, tensor logits of the discriminator)
    """
    # TODO: Implement Function

    ## scope here

    with tf.variable_scope('discriminator', reuse=reuse):

        alpha = 0.2  ### leak relu coeff

        # drop out probability
        keep_prob = 0.8

        # input layer 28 * 28 * color channel
        x1 = tf.layers.conv2d(images, 128, 5, strides=2, padding='same',
                              kernel_initializer= tf.contrib.layers.xavier_initializer(seed=2))
        ## No batch norm here
        ## leak relu here / alpha = 0.2
        relu1 = tf.maximum(alpha * x1, x1)
        # applied drop out here
        drop1 = tf.nn.dropout(relu1, keep_prob= keep_prob)
        # 14 * 14 * 128

        # Layer 2
        x2 = tf.layers.conv2d(drop1, 256, 5, strides=2, padding='same',
                             kernel_initializer= tf.contrib.layers.xavier_initializer(seed=2))
        ## employ batch norm here
        bn2 = tf.layers.batch_normalization(x2, training=True)
        ## leak relu
        relu2 = tf.maximum(alpha * bn2, bn2)
        drop2 = tf.nn.dropout(relu2, keep_prob=keep_prob)

        # 7 * 7 * 256

        # Layer3
        x3 = tf.layers.conv2d(drop2, 512, 5, strides=2, padding='same',
                             kernel_initializer= tf.contrib.layers.xavier_initializer(seed=2))
        bn3 = tf.layers.batch_normalization(x3, training=True)
        relu3 = tf.maximum(alpha * bn3, bn3)
        drop3 = tf.nn.dropout(relu3, keep_prob=keep_prob)
        # 4 * 4 * 512

        # Output
        # Flatten
        flatten = tf.reshape(relu3, (-1, 4 * 4 * 512))
        logits = tf.layers.dense(flatten,1)
        # activation
        out = tf.nn.sigmoid(logits)

    return out, logits

def generator(z, out_channel_dim, is_train=True):
    """
    Create the generator network
    :param z: Input z
    :param out_channel_dim: The number of channels in the output image
    :param is_train: Boolean if generator is being used for training
    :return: The tensor output of the generator
    """
    # TODO: Implement Function

    with tf.variable_scope('generator', reuse = not is_train):
        # First Fully connect layer
        x0 = tf.layers.dense(z, 4 * 4 * 512)
        # Reshape
        x0 = tf.reshape(x0,(-1,4,4,512))
        # Use the batch norm
        bn0 = tf.layers.batch_normalization(x0, training= is_train)
        # Leak relu
        relu0 = tf.nn.relu(bn0)
        # 4 * 4 * 512

        # Conv transpose here
        x1 = tf.layers.conv2d_transpose(relu0, 256, 4, strides=1, padding='valid')
        bn1 = tf.layers.batch_normalization(x1, training=is_train)
        relu1 = tf.nn.relu(bn1)
        # 7 * 7 * 256

        x2 = tf.layers.conv2d_transpose(relu1, 128, 3, strides=2, padding='same')
        bn2 = tf.layers.batch_normalization(x2, training=is_train)
        relu2 = tf.nn.relu(bn2)
        # 14 * 14 * 128

        # Last cov
        logits = tf.layers.conv2d_transpose(relu2, out_channel_dim, 3, strides=2, padding='same')
        ## without batch norm here
        out = tf.tanh(logits)


        return out

def model_loss(input_real, input_z, out_channel_dim):
    """
    Get the loss for the discriminator and generator
    :param input_real: Images from the real dataset
    :param input_z: Z input
    :param out_channel_dim: The number of channels in the output image
    :return: A tuple of (discriminator loss, generator loss)
    """
    # TODO: Implement Function


    g_model = generator(input_z, out_channel_dim, is_train=True)

    d_model_real, d_logits_real = discriminator(input_real, reuse = False)

    d_model_fake, d_logits_fake = discriminator(g_model, reuse= True)

    ## add smooth here

    smooth = 0.1
    d_loss_real = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real,
                                                labels=tf.ones_like(d_model_real) * (1 - smooth)))

    d_loss_fake = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.zeros_like(d_model_fake)))

    g_loss = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake,
                                                labels= tf.ones_like(d_model_fake)))

    d_loss = d_loss_real + d_loss_fake



    return d_loss, g_loss

def model_opt(d_loss, g_loss, learning_rate, beta1):
    """
    Get optimization operations
    :param d_loss: Discriminator loss Tensor
    :param g_loss: Generator loss Tensor
    :param learning_rate: Learning Rate Placeholder
    :param beta1: The exponential decay rate for the 1st moment in the optimizer
    :return: A tuple of (discriminator training operation, generator training operation)
    """

    t_vars = tf.trainable_variables()
    d_vars = [var for var in t_vars if var.name.startswith('discriminator')]
    g_vars = [var for var in t_vars if var.name.startswith('generator')]



    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

    with tf.control_dependencies(update_ops):
        d_train_opt = tf.train.AdamOptimizer(learning_rate,beta1=beta1).minimize(d_loss,var_list = d_vars)
        g_train_opt = tf.train.AdamOptimizer(learning_rate,beta1=beta1).minimize(g_loss,var_list = g_vars)

    return d_train_opt, g_train_opt

def train(epoch_count, batch_size, z_dim, learning_rate, beta1, get_batches, data_shape, data_image_mode):
    """
    Train the GAN
    :param epoch_count: Number of epochs
    :param batch_size: Batch Size
    :param z_dim: Z dimension
    :param learning_rate: Learning Rate
    :param beta1: The exponential decay rate for the 1st moment in the optimizer
    :param get_batches: Function to get batches
    :param data_shape: Shape of the data
    :param data_image_mode: The image mode to use for images ("RGB" or "L")
    """
    losses = []
    samples = []

    input_real, input_z, lr = model_inputs(data_shape[1], data_shape[2], data_shape[3], z_dim)

    d_loss, g_loss = model_loss(input_real,input_z,data_shape[-1])

    d_opt, g_opt = model_opt(d_loss, g_loss, learning_rate, beta1)

    steps = 0

    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        for epoch_i in range(epoch_count):
            for batch_images in get_batches(batch_size):
                # TODO: Train Model
                steps += 1

                # Reshape the image and pass to Discriminator
                batch_images = batch_images.reshape(batch_size,
                                                    data_shape[1],
                                                    data_shape[2],
                                                    data_shape[3])
                # Rescale the data to -1 and 1
                batch_images = batch_images * 2

                # Sample the noise
                batch_z = np.random.uniform(-1,1,size = (batch_size, z_dim))


                ## Run optimizer
                _ = sess.run(d_opt, feed_dict = {input_real:batch_images,
                                                 input_z:batch_z,
                                                 lr:learning_rate
                                                 })
                _ = sess.run(g_opt, feed_dict = {input_real:batch_images,
                                                 input_z:batch_z,
                                                 lr:learning_rate})

                if steps % 10 == 0:

                    train_loss_d = d_loss.eval({input_real:batch_images, input_z:batch_z})
                    train_loss_g = g_loss.eval({input_real:batch_images, input_z:batch_z})

                    losses.append((train_loss_d,train_loss_g))

                    print("Epoch {}/{}...".format(epoch_i+1, epochs),
                          "Discriminator Loss: {:.4f}...".format(train_loss_d),
                          "Generator Loss: {:.4f}".format(train_loss_g))

                if steps % 100 == 0:

                    show_generator_output(sess, 25, input_z, data_shape[-1], data_image_mode)

batch_size = 64
z_dim = 100
learning_rate = 0.001
beta1 = 0.5
epochs = 2

mnist_dataset = helper.Dataset('mnist', glob(os.path.join(data_dir, 'mnist/*.jpg')))
with tf.Graph().as_default():
    train(epochs, batch_size, z_dim, learning_rate, beta1, mnist_dataset.get_batches,
          mnist_dataset.shape, mnist_dataset.image_mode)

batch_size = 32
z_dim = 100
learning_rate = 0.001
beta1 = 0.4
epochs = 1

celeba_dataset = helper.Dataset('celeba', glob(os.path.join(data_dir, 'img_align_celeba/*.jpg')))
with tf.Graph().as_default():
    train(epochs, batch_size, z_dim, learning_rate, beta1, celeba_dataset.get_batches,
          celeba_dataset.shape, celeba_dataset.image_mode)