Generative Adversarial Networks

from __future__ import absolute_import, division, print_function, unicode_literals
from tensorflow.keras import layers
from IPython import display
import tensorflow as tf
import glob
import imageio
import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import time

#数据数量
BUFFER_SIZE = 2656
#批大小
BATCH_SIZE = 16
#迭代次数
EPOCHS = 800
#G的测试噪音大小
noise_dim = 100
#G的测试噪音集大小
num_example_to_G = 4
#dropout_rate
DROPOUT = 0.3
#G的测试集
seed = tf.random.normal([num_example_to_G, noise_dim])



#导入数据集
data = np.load('./paints200.npy')
data = data.astype('float32')
data = (data - 127.5) / 127.5

train_dataset = tf.data.Dataset.from_tensor_slices(data).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)




#生成器与判别器
def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(25*25*256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((25, 25, 256)))
    assert model.output_shape == (None, 25, 25, 256) 

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 50, 50, 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, 100, 100, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

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

    return model

def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=[200, 200, 3]))

    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(DROPOUT))

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

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

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

    return model
generator = tf.keras.models.load_model('epoch7_Gmodel.h5')
discriminator = tf.keras.models.load_model('epoch7_Dmodel.h5')




#损失值
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
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

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




#优化器
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)





#训练
@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)

            display.clear_output(wait=True)
            generate_and_save_images(generator, epoch+1, seed)
            """
            if (epoch + 1) % 100 == 0:
                generator.save('epoch'+epoch+'_Gmodel.h5')
                discriminator.save('epoch'+epoch+'_Dmodel.h5')
                """

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

        generator.save('epoch'+str(epoch+1)+'_Gmodel.h5')
        discriminator.save('epoch'+str(epoch+1)+'_Dmodel.h5')

        display.clear_output(wait=True)
        generate_and_save_images(generator, epochs, seed)

def generate_and_save_images(model, epoch, test_input):
    predictions = model(test_input, training=False)

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

    for i in range(predictions.shape[0]):
        plt.subplot(2, 2, i+1)
        plt.imshow((predictions[i, :, :, :]*127.5+127.5)/255)
        plt.axis('off')

    if not os.path.exists('picture'):
        os.mkdir('picture')

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



train(train_dataset, EPOCHS)
generator.save('newGmodel.h5')
discriminator.save('newDmodel.h5')