def test_get_train_estimator_spec(self):
with tf.Graph().as_default():
gan_model = get_dummy_gan_model()
gan_loss = tfgan.gan_loss(gan_model, dummy_loss_fn, dummy_loss_fn)
spec = get_train_estimator_spec(
gan_model,
gan_loss,
Optimizers(self._generator_optimizer, self._discriminator_optimizer),
get_hooks_fn=None, # use default.
is_chief=True)
self.assertEqual(tf.estimator.ModeKeys.TRAIN, spec.mode)
self.assertShapeEqual(np.array(0), spec.loss) # must be a scalar
self.assertIsNotNone(spec.train_op)
self.assertIsNotNone(spec.training_hooks)
def test_get_sync_estimator_spec(self):
"""Make sure spec is loaded with sync hooks for sync opts."""
with tf.Graph().as_default():
gan_model = get_dummy_gan_model()
gan_loss = tfgan.gan_loss(gan_model, dummy_loss_fn, dummy_loss_fn)
g_opt = get_sync_optimizer()
d_opt = get_sync_optimizer()
spec = get_train_estimator_spec(gan_model,
gan_loss,
Optimizers(g_opt, d_opt),
get_hooks_fn=None) # use default.
self.assertLen(spec.training_hooks, 4)
sync_opts = [
hook._sync_optimizer for hook in spec.training_hooks
if isinstance(hook, get_sync_optimizer_hook_type())
]
self.assertLen(sync_opts, 2)
self.assertSetEqual(frozenset(sync_opts), frozenset(
(g_opt, d_opt)))
def estimator_model_fn(cls,
hparams,
features,
labels,
mode,
config=None,
params=None,
decode_hparams=None,
use_tpu=False):
if mode not in [
model_fn_lib.ModeKeys.TRAIN, model_fn_lib.ModeKeys.EVAL,
model_fn_lib.ModeKeys.PREDICT
]:
raise ValueError('Mode not recognized: %s' % mode)
if mode is model_fn_lib.ModeKeys.TRAIN:
is_training = True
else:
is_training = False
hparams = hparams_lib.copy_hparams(hparams)
# Instantiate model
data_parallelism = None
if not use_tpu and config:
data_parallelism = config.data_parallelism
reuse = tf.get_variable_scope().reuse
# Instantiate model
self = cls(hparams,
mode,
data_parallelism=data_parallelism,
decode_hparams=decode_hparams,
_reuse=reuse)
generator_inputs = self.sample_noise()
# rename inputs for clarity
real_data = features['inputs']
img_shape = common_layers.shape_list(real_data)[1:4]
real_data.set_shape([hparams.batch_size] + img_shape)
# To satify the TFGAN API setting real data to none on predict
if mode == tf.estimator.ModeKeys.PREDICT:
real_data = None
optimizers = Optimizers(
tf.compat.v1.train.AdamOptimizer(hparams.generator_lr,
hparams.beta1),
tf.compat.v1.train.AdamOptimizer(hparams.discriminator_lr,
hparams.beta1))
# Creates tfhub modules for both generator and discriminator
def make_discriminator_spec():
input_layer = tf.placeholder(tf.float32, shape=[None] + img_shape)
disc_output = self.discriminator(input_layer, None, mode)
hub.add_signature(inputs=input_layer, outputs=disc_output)
disc_spec = hub.create_module_spec(make_discriminator_spec)
def make_generator_spec():
input_layer = tf.placeholder(
tf.float32,
shape=[None] + common_layers.shape_list(generator_inputs)[1:])
gen_output = self.generator(input_layer, mode)
hub.add_signature(inputs=input_layer, outputs=gen_output)
gen_spec = hub.create_module_spec(make_generator_spec)
# Create the modules
discriminator_module = hub.Module(disc_spec,
name="Discriminator_Module",
trainable=True)
generator_module = hub.Module(gen_spec,
name="Generator_Module",
trainable=True)
# Wraps the modules into functions expected by TF-GAN
def generator(code, mode):
p = hparams
out = generator_module(code)
shape = common_layers.shape_list(out)
# Applying convolution by PSF convolution
if p.apply_psf and 'psf' in features:
out = convolve(out,
tf.cast(features['psf'][..., 0], tf.complex64))
# Adds noise according to the provided power spectrum
noise = tf.spectral.rfft2d(tf.random_normal(out.get_shape()[:3]))
thresholded_ps = tf.where(features['ps'] >= 9,
tf.zeros_like(features['ps']),
tf.sqrt(tf.exp(features['ps'])))
noise = noise * tf.cast(thresholded_ps, tf.complex64)
out = out + tf.expand_dims(tf.spectral.irfft2d(noise), axis=-1)
return out
discriminator = lambda image, conditioning, mode: discriminator_module(
image)
# Make GANModel, which encapsulates the GAN model architectures.
gan_model = get_gan_model(mode,
generator,
discriminator,
real_data,
generator_inputs,
add_summaries=self.summaries)
# Make GANLoss, which encapsulates the losses.
if mode in [tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL]:
gan_loss = tfgan_train.gan_loss(gan_model,
self.generator_loss,
self.discriminator_loss,
add_summaries=True)
# Make the EstimatorSpec, which incorporates the GANModel, losses, eval
# metrics, and optimizers (if required).
if mode == tf.estimator.ModeKeys.TRAIN:
get_hooks_fn = tfgan_train.get_sequential_train_hooks(
namedtuples.GANTrainSteps(hparams.gen_steps,
hparams.disc_steps))
estimator_spec = get_train_estimator_spec(gan_model,
gan_loss,
optimizers,
get_hooks_fn,
is_chief=True)
elif mode == tf.estimator.ModeKeys.EVAL:
estimator_spec = get_eval_estimator_spec(gan_model, gan_loss)
else: # tf.estimator.ModeKeys.PREDICT
# Register hub modules for export
hub.register_module_for_export(generator_module, "generator")
hub.register_module_for_export(discriminator_module,
"discriminator")
estimator_spec = get_predict_estimator_spec(gan_model)
return estimator_spec
