class Vanilla_gan():
def __init__(self, config, summary):
self.summary = summary
# ...
# Build the generator and critic
self.generator = self.make_generator_model()
self.discriminator = self.make_discriminator_model()
noise = tf.random.normal([1, 100])
generated_image = self.generator(noise, training=False)
#
print('generated_image.shape ', generated_image.shape)
print('generated_image_min ', np.min((generated_image[0, :, :,:] + 1) / 2))
print('generated_image_max ', np.max((generated_image[0, :, :,:] + 1) / 2))
plt.imshow((generated_image[0, :, :, :] + 1) / 2)
plt.show()
#
decision = self.discriminator(generated_image)
print("decision:", decision)
# This method returns a helper function to compute cross entropy loss
self.cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
self.generator_optimizer = tf.keras.optimizers.Adam(1e-4)
self.discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)
self.summary.summary_writer.set_value('generator', value=self.generator.to_json())
self.summary.summary_writer.update_log()
self.summary.summary_writer.set_value('discriminator', value=self.discriminator.to_json())
self.summary.summary_writer.update_log()
set_up_checkpoint(self)
global SEED
SEED = tf.random.normal([num_examples_to_generate, noise_dim])
#3 Generator/Decoder
def make_generator_model(self):
model = tf.keras.Sequential()
model.add(layers.Dense(units=8*8*256, use_bias=False, input_shape=(100,)))
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Reshape((8, 8, 256)))
assert model.output_shape == (None, 8, 8, 256) # Note: None is the batch size
model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
assert model.output_shape == (None, 8, 8, 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, 16, 16, 64)
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Conv2DTranspose(64, (5, 5), strides=(3, 3), padding='same', use_bias=False))
assert model.output_shape == (None, 48, 48, 64)
model.add(layers.BatchNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Conv2DTranspose(3, (5, 5), strides=(3, 3), padding='same', use_bias=False, activation='tanh'))
assert model.output_shape == (None, 144, 144, 3)
return model
#4 Discriminator
def make_discriminator_model(self):
model = tf.keras.Sequential()
model.add(layers.Conv2D(filters = 64, kernel_size = (5, 5), strides=(2, 2), padding='same', input_shape=[TILE_SIZE, TILE_SIZE, 3]))
# model.add(layers.LayerNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dropout(0.3))
model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
# model.add(layers.LayerNormalization())
model.add(layers.LeakyReLU())
model.add(layers.Dropout(0.3))
model.add(layers.Flatten())
model.add(layers.Dense(1))
return model
#5 Losses
def discriminator_loss(self, real_output, fake_output):
real_loss = self.cross_entropy(tf.ones_like(real_output), real_output)
fake_loss = self.cross_entropy(tf.zeros_like(fake_output), fake_output)
# should not be total_loss = (real_loss + fake_loss) * 0.5 ?
total_loss = real_loss + fake_loss
return total_loss
def generator_loss(self, fake_output):
return self.cross_entropy(tf.ones_like(fake_output), fake_output)
#6 Train def
@tf.function
def train_step_generator(self, images):
noise = tf.random.normal([BATCH_SIZE, noise_dim])
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
generated_images = self.generator(noise, training=True)
fake_output = self.discriminator(generated_images, training=False)
gen_loss = self.generator_loss(fake_output)
gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
return gen_loss
return gen_loss
@tf.function
def train_step_discriminator(self, images):
# print('images.shape: ', images.shape)
noise = tf.random.normal([BATCH_SIZE, noise_dim])
# noise = tf.random.normal([BATCH_SIZE, noise_dim])
with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
generated_images = self.generator(noise, training=False)
real_output = self.discriminator(images, training=True)
fake_output = self.discriminator(generated_images, training=True)
disc_loss = self.discriminator_loss(real_output, fake_output)
gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)
self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))
return disc_loss
def train(self, dataset, epochs):
start = time.time()
# statistics of loss functions
gen_loss, disc_loss, gen_loss_stat, disc_loss_stat = [], [], [], []
for epoch in range(epochs):
epoch_start = time.time()
# dataset[epoch][0] - images
# dataset[epoch][1] - labels
for batch_num in range(int(1024/16)):
# train_gen.unaugmented(0)[0].shape
# (BATCH_SIZE, 9, 96, 96, 3)
batch = dataset.unaugmented(batch_num)[0]
disc_loss_sample = self.train_step_discriminator(batch)
gen_loss_sample =1 self.train_step_generator(batch)
gen_loss.append(gen_loss_sample)
disc_loss.append(disc_loss_sample)
# print('loss: disc: ', disc_loss_sample, 'gen: ', gen_loss_sample)
g_loss = sum(gen_loss)/len(gen_loss)
d_loss = sum(disc_loss)/len(disc_loss)
epoch_elapsed = time.time()-epoch_start
print (f'Epoch {epoch+1}, gen loss={g_loss},disc loss={d_loss}, {hms_string(epoch_elapsed)}')
gen_loss_stat.append(g_loss)
disc_loss_stat.append(d_loss)
# Log
# d_loss_list = list(map(lambda x: str(x.item()), d_loss))
self.summary.summary_writer.add_value('loss', value=[epoch, str(d_loss.numpy()), str(g_loss.numpy())])
self.summary.summary_writer.update_log()
# save images and checkpoints
if (epoch + 1) % config['sample_interval'] == 0:
sample_images(self, epoch, SEED, num_examples_to_generate)
if (epoch +1) % config['sample_interval'] == 0:
plot_loss(self, [gen_loss_stat, disc_loss_stat], ["Generator", "Discriminator"], epoch)
self.checkpoint.step.assign_add(1)
if (epoch + 1) % config['save_interval'] == 0:
save_path = self.checkpoint_manager.save()
print("Saved checkpoint for step {}: {}".format(int(self.checkpoint.step), save_path))
# final images
sample_images(self, epochs, SEED, num_examples_to_generate)
return gen_loss_stat, disc_loss_stat