def train(hparams):
"""Trains a StarGAN.
Args:
hparams: An HParams instance containing the hyperparameters for training.
"""
# Create the log_dir if not exist.
if not tf.io.gfile.exists(hparams.train_log_dir):
tf.io.gfile.makedirs(hparams.train_log_dir)
# Shard the model to different parameter servers.
with tf.device(tf.compat.v1.train.replica_device_setter(hparams.ps_replicas)):
# Create the input dataset.
with tf.compat.v1.name_scope('inputs'), tf.device('/cpu:0'):
images, labels = data_provider.provide_data('train', hparams.batch_size,
hparams.patch_size)
# Define the model.
with tf.compat.v1.name_scope('model'):
model = _define_model(images, labels)
# Add image summary.
tfgan.eval.add_stargan_image_summaries(
model, num_images=3 * hparams.batch_size, display_diffs=True)
# Define the model loss.
loss = tfgan.stargan_loss(model)
# Define the train ops.
with tf.compat.v1.name_scope('train_ops'):
train_ops = _define_train_ops(model, loss, hparams.generator_lr,
hparams.discriminator_lr,
hparams.adam_beta1, hparams.adam_beta2,
hparams.max_number_of_steps)
# Define the train steps.
train_steps = _define_train_step(hparams.gen_disc_step_ratio)
# Define a status message.
status_message = tf.strings.join([
'Starting train step: ',
tf.as_string(tf.compat.v1.train.get_or_create_global_step())
],
name='status_message')
# Train the model.
tfgan.gan_train(
train_ops,
hparams.train_log_dir,
get_hooks_fn=tfgan.get_sequential_train_hooks(train_steps),
hooks=[
tf.estimator.StopAtStepHook(num_steps=hparams.max_number_of_steps),
tf.estimator.LoggingTensorHook([status_message], every_n_iter=10)
],
master=hparams.tf_master,
is_chief=hparams.task == 0)
def test_provide_data(self, mock_tfds):
batch_size = 5
patch_size = 32
mock_tfds.load.return_value = self.mock_ds
images, labels = data_provider.provide_data(
'test', batch_size, patch_size=patch_size, domains=('A', 'B', 'C'))
self.assertLen(images, 3)
self.assertLen(labels, 3)
with self.cached_session() as sess:
images = sess.run(images)
labels = sess.run(labels)
for img in images:
self.assertTupleEqual(img.shape, (batch_size, patch_size, patch_size, 3))
self.assertTrue(np.all(np.abs(img) <= 1))
for lbl in labels:
expected_lbls_shape = (batch_size, 3)
self.assertTupleEqual(lbl.shape, expected_lbls_shape)