Fix wasserstein gradient penalty name scope issue and add the proper …

…name scope. PiperOrigin-RevId: 171610946
tensorflow · caisq · Oct 6, 2017 · Oct 6, 2017 · Oct 6, 2017 · Oct 6, 2017
commit 52d3a842463d11990600bb65f9752b59f6d8f418
diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl.py b/tensorflow/contrib/gan/python/losses/python/losses_impl.py
@@ -297,7 +297,6 @@ def acgan_generator_loss(
 # GANs` (https://arxiv.org/abs/1704.00028).
 
 
-# TODO(joelshor): Figure out why this function can't be inside a name scope.
 def wasserstein_gradient_penalty(
     real_data,
     generated_data,
@@ -339,48 +338,50 @@ def wasserstein_gradient_penalty(
   Raises:
     ValueError: If the rank of data Tensors is unknown.
   """
-  real_data = ops.convert_to_tensor(real_data)
-  generated_data = ops.convert_to_tensor(generated_data)
-  if real_data.shape.ndims is None:
-    raise ValueError('`real_data` can\'t have unknown rank.')
-  if generated_data.shape.ndims is None:
-    raise ValueError('`generated_data` can\'t have unknown rank.')
-
-  differences = generated_data - real_data
-  batch_size = differences.shape[0].value or array_ops.shape(differences)[0]
-  alpha_shape = [batch_size] + [1] * (differences.shape.ndims - 1)
-  alpha = random_ops.random_uniform(shape=alpha_shape)
-  interpolates = real_data + (alpha * differences)
-
-  # Reuse variables if a discriminator scope already exists.
-  reuse = False if discriminator_scope is None else True
-  with variable_scope.variable_scope(discriminator_scope, 'gpenalty_dscope',
-                                     reuse=reuse):
-    disc_interpolates = discriminator_fn(interpolates, generator_inputs)
-
-  if isinstance(disc_interpolates, tuple):
-    # ACGAN case: disc outputs more than one tensor
-    disc_interpolates = disc_interpolates[0]
-
-  gradients = gradients_impl.gradients(disc_interpolates, interpolates)[0]
-  gradient_squares = math_ops.reduce_sum(
-      math_ops.square(gradients), axis=list(range(1, gradients.shape.ndims)))
-  # Propagate shape information, if possible.
-  if isinstance(batch_size, int):
-    gradient_squares.set_shape([
-        batch_size] + gradient_squares.shape.as_list()[1:])
-  # For numerical stability, add epsilon to the sum before taking the square
-  # root. Note tf.norm does not add epsilon.
-  slopes = math_ops.sqrt(gradient_squares + epsilon)
-  penalties = math_ops.square(slopes - 1.0)
-  penalty = losses.compute_weighted_loss(
-      penalties, weights, scope=scope, loss_collection=loss_collection,
-      reduction=reduction)
+  with ops.name_scope(scope, 'wasserstein_gradient_penalty',
+                      (real_data, generated_data)) as scope:
+    real_data = ops.convert_to_tensor(real_data)
+    generated_data = ops.convert_to_tensor(generated_data)
+    if real_data.shape.ndims is None:
+      raise ValueError('`real_data` can\'t have unknown rank.')
+    if generated_data.shape.ndims is None:
+      raise ValueError('`generated_data` can\'t have unknown rank.')
+
+    differences = generated_data - real_data
+    batch_size = differences.shape[0].value or array_ops.shape(differences)[0]
+    alpha_shape = [batch_size] + [1] * (differences.shape.ndims - 1)
+    alpha = random_ops.random_uniform(shape=alpha_shape)
+    interpolates = real_data + (alpha * differences)
+
+    with ops.name_scope(None):  # Clear scope so update ops are added properly.
+      # Reuse variables if variables already exists.
+      with variable_scope.variable_scope(discriminator_scope, 'gpenalty_dscope',
+                                         reuse=variable_scope.AUTO_REUSE):
+        disc_interpolates = discriminator_fn(interpolates, generator_inputs)
+
+    if isinstance(disc_interpolates, tuple):
+      # ACGAN case: disc outputs more than one tensor
+      disc_interpolates = disc_interpolates[0]
+
+    gradients = gradients_impl.gradients(disc_interpolates, interpolates)[0]
+    gradient_squares = math_ops.reduce_sum(
+        math_ops.square(gradients), axis=list(range(1, gradients.shape.ndims)))
+    # Propagate shape information, if possible.
+    if isinstance(batch_size, int):
+      gradient_squares.set_shape([
+          batch_size] + gradient_squares.shape.as_list()[1:])
+    # For numerical stability, add epsilon to the sum before taking the square
+    # root. Note tf.norm does not add epsilon.
+    slopes = math_ops.sqrt(gradient_squares + epsilon)
+    penalties = math_ops.square(slopes - 1.0)
+    penalty = losses.compute_weighted_loss(
+        penalties, weights, scope=scope, loss_collection=loss_collection,
+        reduction=reduction)
 
-  if add_summaries:
-    summary.scalar('gradient_penalty_loss', penalty)
+    if add_summaries:
+      summary.scalar('gradient_penalty_loss', penalty)
 
-  return penalty
+    return penalty
 
 
 # Original losses from `Generative Adversarial Nets`

diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py b/tensorflow/contrib/gan/python/losses/python/losses_impl_test.py
@@ -453,10 +453,11 @@ def setUp(self):
         'discriminator_scope': self._scope,
     }
     self._expected_loss = 9.00000
-    self._expected_op_name = 'weighted_loss/value'
+    self._expected_op_name = 'wasserstein_gradient_penalty/value'
     self._batch_size = 1
 
   def _discriminator_fn(self, inputs, _):
+    ops.add_to_collection('fake_update_ops', constant_op.constant(1.0))
     return variable_scope.get_variable('dummy_d', initializer=2.0) * inputs
 
   def test_loss_with_placeholder(self):
@@ -487,6 +488,26 @@ def test_reuses_scope(self):
     self.assertEqual(
         num_vars, len(ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)))
 
+  def test_works_with_get_collection(self):
+    """Tests that gradient penalty works inside other scopes."""
+    # We ran the discriminator once in the setup, so there should be an op
+    # already in the collection.
+    self.assertEqual(1, len(ops.get_collection(
+        'fake_update_ops', self._kwargs['discriminator_scope'].name)))
+
+    # Make sure the op is added to the collection even if it's in a name scope.
+    with ops.name_scope('loss'):
+      tfgan_losses.wasserstein_gradient_penalty(**self._kwargs)
+    self.assertEqual(2, len(ops.get_collection(
+        'fake_update_ops', self._kwargs['discriminator_scope'].name)))
+
+    # Make sure the op is added to the collection even if it's in a variable
+    # scope.
+    with variable_scope.variable_scope('loss_vscope'):
+      tfgan_losses.wasserstein_gradient_penalty(**self._kwargs)
+    self.assertEqual(3, len(ops.get_collection(
+        'fake_update_ops', self._kwargs['discriminator_scope'].name)))
+
 
 class MutualInformationPenaltyTest(test.TestCase, _PenaltyTest):
   """Tests for mutual_information_penalty."""