diff options
| author |  Wesley Qian <wwq@google.com> | 2018-08-13 10:01:45 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-13 10:06:26 -0700 |
| commit | bc89890f52a00b6094f2c49ec76a7264680beb8c (patch) | |
| tree | 37a178d020a82c5838160ddaece0c262a3e69ef3 /tensorflow/contrib/gan | |
| parent | ec3d9d9cf32697b0a44c1738dd06ec79f8db9064 (diff) | |
Expose stargan_loss in train.py.
PiperOrigin-RevId: 208497573
Diffstat (limited to 'tensorflow/contrib/gan')
| -rw-r--r-- | tensorflow/contrib/gan/python/train.py | 14 |
1 files changed, 8 insertions, 6 deletions
diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index 03f52d214b..9e5aea1498 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -52,7 +52,6 @@ from tensorflow.python.training import session_run_hook from tensorflow.python.training import sync_replicas_optimizer from tensorflow.python.training import training_util - __all__ = [ 'gan_model', 'infogan_model', @@ -61,6 +60,7 @@ __all__ = [ 'stargan_model', 'gan_loss', 'cyclegan_loss', + 'stargan_loss', 'gan_train_ops', 'gan_train', 'get_sequential_train_hooks', @@ -646,8 +646,9 @@ def gan_loss( type(model)) # Optionally create pooled model. - pooled_model = (_tensor_pool_adjusted_model(model, tensor_pool_fn) if - tensor_pool_fn else model) + pooled_model = ( + _tensor_pool_adjusted_model(model, tensor_pool_fn) + if tensor_pool_fn else model) # Create standard losses. gen_loss = generator_loss_fn(model, add_summaries=add_summaries) @@ -665,9 +666,10 @@ def gan_loss( if _use_aux_loss(mutual_information_penalty_weight): gen_info_loss = tfgan_losses.mutual_information_penalty( model, add_summaries=add_summaries) - dis_info_loss = (gen_info_loss if tensor_pool_fn is None else - tfgan_losses.mutual_information_penalty( - pooled_model, add_summaries=add_summaries)) + dis_info_loss = ( + gen_info_loss + if tensor_pool_fn is None else tfgan_losses.mutual_information_penalty( + pooled_model, add_summaries=add_summaries)) gen_loss += mutual_information_penalty_weight * gen_info_loss dis_loss += mutual_information_penalty_weight * dis_info_loss if _use_aux_loss(aux_cond_generator_weight): |