def test_get_gan_model(self, mode): with tf.Graph().as_default(): generator_inputs = {'x': tf.ones([3, 4])} is_predict = mode == tf.estimator.ModeKeys.PREDICT real_data = tf.zeros([3, 4]) if not is_predict else None gan_model = get_gan_model(mode, generator_fn, discriminator_fn, real_data, generator_inputs, add_summaries=False) self.assertEqual(generator_inputs, gan_model.generator_inputs) self.assertIsNotNone(gan_model.generated_data) self.assertLen(gan_model.generator_variables, 2) # 1 FC layer self.assertIsNotNone(gan_model.generator_fn) if mode == tf.estimator.ModeKeys.PREDICT: self.assertIsNone(gan_model.real_data) self.assertIsNone(gan_model.discriminator_real_outputs) self.assertIsNone(gan_model.discriminator_gen_outputs) self.assertIsNone(gan_model.discriminator_variables) self.assertIsNone(gan_model.discriminator_scope) self.assertIsNone(gan_model.discriminator_fn) else: self.assertIsNotNone(gan_model.real_data) self.assertIsNotNone(gan_model.discriminator_real_outputs) self.assertIsNotNone(gan_model.discriminator_gen_outputs) self.assertLen(gan_model.discriminator_variables, 2) # 1 FC layer self.assertIsNotNone(gan_model.discriminator_scope) self.assertIsNotNone(gan_model.discriminator_fn)
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