def run(train_data, valid_data, len_size, scale, EPOCHS, root_path='./', load_model_dir=None, saved_model_dir=None, log_dir=None, summary=False):
if log_dir is None:
log_dir = os.path.join(root_path, 'our_model', 'logs', 'model')
logging.info(train_data)
logging.info(valid_data)
# get generator model and discriminator model
Gen = model.make_generator_model(len_high_size=len_size, scale=scale)
Dis = model.make_discriminator_model(len_high_size=len_size, scale=scale)
if load_model_dir is not None:
#load_model_dir = os.path.join(root_path, 'our_model', 'saved_model')
file_path = os.path.join(load_model_dir, 'gen_model_'+str(len_size), 'gen_weights')
if os.path.exists(file_path):
Gen.load_weights(file_path)
else:
logging.info("generator doesn't exist. create a new one.")
file_path = os.path.join(load_model_dir, 'dis_model_'+str(len_size), 'dis_weights')
if os.path.exists(file_path):
Dis.load_weights(file_path)
else:
logging.info("discriminator model doesn't exist. create a new one")
if summary:
logging.info(Gen.summary())
tf.keras.utils.plot_model(Gen, to_file='G.png', show_shapes=True)
logging.info(Dis.summary())
tf.keras.utils.plot_model(Dis, to_file='D.png', show_shapes=True)
if saved_model_dir is None:
saved_model_dir = os.path.join(root_path, 'our_model', 'saved_model')
model.fit(Gen, Dis, train_data, EPOCHS, len_size, scale,
valid_data, log_dir=log_dir, saved_model_dir=saved_model_dir)
file_path = os.path.join(
saved_model_dir, 'gen_model_'+str(len_size), 'gen_weights')
Gen.save_weights(file_path)
file_path = os.path.join(
saved_model_dir, 'dis_model_'+str(len_size), 'dis_weights')
Dis.save_weights(file_path)
def main(input_array, input_width=16, input_height=22, output_width=100, output_height=64):
input_array = (input_array - 127.5) / 127.5
train_dataset = tf.data.Dataset.from_tensor_slices(input_array).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
generator = make_generator_model(input_width, input_height, output_width, output_height)
discriminator = make_discriminator_model(output_width, output_height)
generator_optimizer = tf.keras.optimizers.Adam(1e-4)
discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
discriminator_optimizer=discriminator_optimizer,
generator=generator,
discriminator=discriminator)
seed = tf.random.normal([num_examples_to_generate, noise_dim])
@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 generate_and_save_images(model, epoch, test_input):
# Notice `training` is set to False.
# This is so all layers run in inference mode (batchnorm).
predictions = model(test_input, training=False)
fig = plt.figure(figsize=(4, 4))
for i in range(predictions.shape[0]):
plt.subplot(4, 4, i + 1)
plt.imshow(predictions[i, :, :, :3] * 127.5 + 127.5)
plt.axis('off')
plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
plt.show()
def train(dataset, epochs):
for epoch in range(epochs):
start = time.time()
for image_batch in dataset:
train_step(image_batch)
# Produce images for the GIF as we go
generate_and_save_images(generator,
epoch + 1,
seed)
# Save the model every 15 epochs
if (epoch + 1) % 20 == 0:
checkpoint.save(file_prefix=checkpoint_prefix)
generator.save_weights("model_epoch_"+ str(epoch) + ".h5")
print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))
# Generate after the final epoch
generate_and_save_images(generator,
epochs,
seed)
train(train_dataset, EPOCHS)
import tensorflow as tf
import os
import time
import numpy as np
from define import *
import model
# build model
generator = model.make_generator_model()
discriminator = model.make_discriminator_model()
# checkpoint
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(generator=generator,
discriminator=discriminator)
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
seeds = []
confs = []
# predict
for epoch in range(100):
# make seed
seed = tf.random.normal([num_examples_to_generate, noise_dim])
start = time.time()
predictions = generator(seed, training=False)