在这篇文章中,我们使用Tensorflow2.0来实现GAN,使用的数据集是手写数字数据集。
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
%matplotlib inline
(train_images, train_labels), (_, _) = keras.datasets.mnist.load_data()
train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')
train_images = (train_images-127.5)/127.5
BATCH_SIZE = 256
BUFFER_SIZE = 60000
datasets = tf.data.Dataset.from_tensor_slices(train_images)
datasets = datasets.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
def generator_model(): # 用100个随机数(噪音)生成手写数据集
model = keras.Sequential()
model.add(layers.Dense(256, input_shape=(100,), use_bias=False))
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dense(512, use_bias=False))
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dense(28*28*1, use_bias=False, activation='tanh'))
model.add(layers.BatchNormalization())
model.add(layers.Reshape((28, 28, 1)))
return model
def discriminator_model(): # 识别输入的图片
model = keras.Sequential()
model.add(layers.Flatten())
model.add(layers.Dense(512, use_bias=False))
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dense(256, use_bias=False))
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dense(1))
return model
对于判别器来说,我们需要将导入的原始图片识别为真(1),将生成器胜场的图像识别为假(0)。
对于生成器来说,我们需要使得生成的图片无限接近于真实图片。
cross_entropy = keras.losses.BinaryCrossentropy(from_logits=True)
def discriminator_loss(real_out, fake_out):
real_loss = cross_entropy(tf.ones_like(real_out), real_out)
fake_loss = cross_entropy(tf.zeros_like(fake_out), fake_out)
return real_loss + fake_loss
def generator_loss(fake_out):
return cross_entropy(tf.ones_like(fake_out), fake_out)
在以上代码中,real_out是指向判别器输入原始图像得到的结果;fake_out是指向判别器输入生成图像得到的结果。
所以对于判别器的损失函数来说,real_out应该无限接近于1;fake_out应该无限接近于0。即我们想训练出的判别器应该对图片有很高的识别能力。
但对于生成器的损失函数来说,fake_out应该无限接近于1,也就是令判别器很难分辨出生成的图片。
【注】keras.losses.BinaryCrossentropy(from_logits=True)的用法可以参考:tensorflow2.0中损失函数tf.keras.losses.BinaryCrossentropy()的用法。
generator_opt = keras.optimizers.Adam(1e-4)
discriminator_opt = keras.optimizers.Adam(1e-4)
generator = generator_model()
discriminator = discriminator_model()
noise_dim = 100 # 即用100个随机数生成图片
def train_step(images):
noise = tf.random.normal([BATCH_SIZE, noise_dim])
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
real_out = discriminator(images, training=True)
gen_image = generator(noise, training=True)
fake_out = discriminator(gen_image, training=True)
gen_loss = generator_loss(fake_out)
disc_loss = discriminator_loss(real_out, fake_out)
gradient_gen = gen_tape.gradient(gen_loss, generator.trainable_variables)
gradient_disc = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
generator_opt.apply_gradients(zip(gradient_gen, generator.trainable_variables))
discriminator_opt.apply_gradients(zip(gradient_disc, discriminator.trainable_variables))
gradient_gen = gen_tape.gradient(gen_loss, generator.trainable_variables)表示计算gen_loss对于generator的所有变量的梯度。
generator_opt.apply_gradients(zip(gradient_gen, generator.trainable_variables))表示根据gradient_gen来优化generator的变量。
【注】梯度带及梯度更新的用法参考:Tensorflow中的梯度带(GradientTape)以及梯度更新。
def generate_plot_image(gen_model, test_noise):
pre_images = gen_model(test_noise, training=False)
fig = plt.figure(figsize=(4, 4))
for i in range(pre_images.shape[0]):
plt.subplot(4, 4, i+1)
plt.imshow((pre_images[i, :, :, 0] + 1)/2, cmap='gray')
plt.axis('off')
plt.show()
plt.imshow((pre_images[i, :, :, 0] + 1)/2, cmap=‘gray’)
这里是因为我们使用tanh激活函数之后会将结果限制在-1到1之间,而我们需要将其转化到0到1之间。
EPOCHS = 100 # 训练100次
num_exp_to_generate = 16 # 生成16张图片
seed = tf.random.normal([num_exp_to_generate, noise_dim]) # 16组随机数组,每组含100个随机数,用来生成16张图片。
def train(dataset, epochs):
for epoch in range(epochs):
for image_batch in dataset:
train_step(image_batch)
print('.', end='')
generate_plot_image(generator, seed)
train(datasets, EPOCHS)