GAN 网络在MINIST数据集上的简单运用

温习了一下lab的内容

用的包

用的tensorflow

import tensorflow as tf
keras = tf.keras
layers = tf.keras.layers
# Include numpy for basic stuff
import numpy as np
# matplotlib allows us to visualise our data.
import matplotlib.pyplot as plt
# Import a library for displaying models
from IPython.display import Image

导入数据

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
# Print the shapes of the array and the unique label 
print(x_train.shape, y_train.shape, np.unique(y_test))
# Pick an example index
example_idx = 0
plt.imshow(x_train[example_idx,...],cmap='gray')
# Print the label
plt.show()

预处理数据

# Divide the image data to put it in the right range and convert to floating point numbers
x_train = x_train.astype("float32") / 255
# Preprocess the test data the same way
x_test = x_test.astype("float32") / 255

# Stick the data together
x_data = np.concatenate([x_train, x_test], axis=0)

# Reshape to (-1, im_height, im_width, 1)
x_data = np.reshape(x_data, (-1, 28, 28,1))

# We're not interested in the labels, but if we were...
# Convert the labels to floating point
y_train = y_train.astype("float32")
y_test = y_test.astype("float32")
y_data = np.concatenate([y_train, y_test], axis=0)

# Define the BATCH_SIZE as a variable as we need it later
BATCH_SIZE = 256
# Prepare the training dataset into batches and shuffle the examples
train_dataset = tf.data.Dataset.from_tensor_slices(x_data)
train_dataset = train_dataset.shuffle(buffer_size=10000).batch(BATCH_SIZE)

构建生成器

最初的图片大小为 [28 * 28] ，28+28=49，所以映射成 [7 * 7]

def make_generator(z_size=32, layer_sizes=[256, 128, 64]):
  """
  z_size specifies the latent variable dimension
  layer_sizes: is a list containing the number of channels at each stage in the generator
  """
  # Note that the Input size relates to the size of each instance
  inputs = keras.Input(shape=(z_size,), name="z")
  net = inputs

  # We first need to increase the size of our Tensor to reshape it into an image and
  # apply some convolutions. Let's map our z into a 7x7xlayers[0].  
  net = layers.Dense(7*7*layer_sizes[0],use_bias=False)(net)

  # The Dense lay produces a new vector, which we can reshape to look like an image.
  # Reshape to the right image size
  net = layers.Reshape((7,7,layer_sizes[0]))(net)

  # Apply batch normalisation followed by leaky relu (a slight variation on the relu activation function)
  net = layers.BatchNormalization(scale=False, center=False)(net)
  net = layers.LeakyReLU()(net)
  # Let's do our first convolution
  net = layers.Conv2D(layer_sizes[1], 5, padding='same', use_bias=False)(net)
  net = layers.BatchNormalization(scale=False, center=False)(net)
  net = layers.LeakyReLU()(net)

  # We want to grow the size of the image, let's do this by applying bilinear upsampling
  net = layers.UpSampling2D(size=(2,2))(net)

  # Followed by convolution, batch norm and relu again
  net = layers.Conv2D(layer_sizes[2], 5, padding='same', use_bias=False)(net)
  net = layers.BatchNormalization(scale=False, center=False)(net)
  net = layers.LeakyReLU()(net)

    # We want to grow the size of the image, let's do this by applying bilinear upsampling
  net = layers.UpSampling2D(size=(2,2))(net)

  # We do one final convolution to produce the final image (use a sigmoid to put the values in the correct range)
  net = layers.Conv2D(1, 5, padding='same', activation='sigmoid')(net)
  

  return keras.Model(inputs=inputs, outputs=net)

# Test the function works - but don't save the result yet
make_generator()

构建判别器

def make_discriminator(layer_sizes=[64, 128]):
  """
  layer_sizes: list containing the number of channels at each stage of the discriminator
  """
  # Note that the Input size relates to the size of each instance
  inputs = keras.Input(shape=(28,28,1), name="image")
  net = inputs

  # We can also define the layer object, and store it in a Python variable and call it later
  # We're more likely to want to do this for layers that contain model weights
  conv1 = layers.Conv2D(filters=layer_sizes[0], kernel_size=5, padding='same', strides=(2,2))
  net = conv1(net)
  # Dropout might help a bit
  #net = layers.Dropout(0.1)(net)
  net = layers.LeakyReLU()(net)

  conv2 = layers.Conv2D(filters=layer_sizes[1], kernel_size=5, padding='same', strides=(2,2))
  net = conv2(net)
  # Dropout might help a bit
  #net = layers.Dropout(0.1)(net)
  net = layers.LeakyReLU()(net)


  # We need to flatten the spatial dimensions before putting it through the dense layers
  net = layers.Flatten()(net)
  dense1 = layers.Dense(units=1)
  # Apply the dense layer
  net = dense1(net)

  return keras.Model(inputs=inputs, outputs=net)

make_discriminator()

构建模型

noise_dim = 64

generator = make_generator(z_size=noise_dim, layer_sizes=[256, 128, 64])

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

#plt.imshow(generated_image, cmap='gray')

# Print the output of the discriminator
discriminator = make_discriminator(layer_sizes=[64, 128])
print(discriminator(generated_image))
plt.imshow(generated_image[0, :, :, 0], cmap='gray')

定义模型损失

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

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)

# Define the optimizers
generator_optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)

定义训练step

# Notice the use of `tf.function`
# This annotation causes the function to be "compiled". Making it faster to call thereafter
@tf.function
def train_step(images):
    noise = tf.random.normal([BATCH_SIZE, noise_dim])
    # Create a pair of gradient tapes
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
      # Generate some fake data from the noise input
      generated_images = generator(noise, training=True)
      # Call the discriminator on the real images
      real_output = discriminator(images, training=True)
      # Call the discriminator on the generated images
      fake_output = discriminator(generated_images, training=True)
      # Calculate the losses
      gen_loss = generator_loss(fake_output)
      disc_loss = discriminator_loss(real_output, fake_output)
    # Calculate the gradient of the losses for the right variables
    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
    # Update the model variables
    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
    # Return the losses
    return gen_loss, disc_loss

开始训练

def train(dataset, epochs):
  for epoch in range(epochs):

    for idx, image_batch in enumerate(dataset):
      losses = train_step(image_batch)
      # Every N btaches, print the losses of that batch and draw some example images
      if idx % 10 == 0:
        print(losses)
        noise = tf.random.normal([3, noise_dim])
        generated_image = generator(noise, training=False)
        plt.subplot(131)
        plt.imshow(generated_image[0, :, :, 0], cmap='gray')
        plt.subplot(132)
        plt.imshow(generated_image[1, :, :, 0], cmap='gray')
        plt.subplot(133)
        plt.imshow(generated_image[2, :, :, 0], cmap='gray')
        plt.show()

# 
train(train_dataset, 30)

运行结果

并不是很理想