code: # coding=utf-8
  # Copyright 2020 The TensorFlow GAN Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # python2 python3
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  import tensorflow as tf
  
  from tensorflow_gan.examples.progressive_gan import layers
  from tensorflow_gan.examples.progressive_gan import networks
  
  
  def _get_grad_norm(ys, xs):
    """Compute 2-norm of dys / dxs."""
    return tf.sqrt(
        tf.add_n([
            tf.reduce_sum(input_tensor=tf.square(g))
            for g in tf.gradients(ys=ys, xs=xs)
        ]))
  
  
  def _num_filters_stub(block_id):
    return networks.num_filters(block_id, 8, 1, 8)
  
  
  class NetworksTest(tf.test.TestCase):
  
    def test_resolution_schedule_correct(self):
      rs = networks.ResolutionSchedule(
          start_resolutions=[5, 3], scale_base=2, num_resolutions=3)
      self.assertEqual(rs.start_resolutions, (5, 3))
      self.assertEqual(rs.scale_base, 2)
      self.assertEqual(rs.num_resolutions, 3)
      self.assertEqual(rs.final_resolutions, (20, 12))
      self.assertEqual(rs.scale_factor(1), 4)
      self.assertEqual(rs.scale_factor(2), 2)
      self.assertEqual(rs.scale_factor(3), 1)
      with self.assertRaises(ValueError):
        rs.scale_factor(0)
      with self.assertRaises(ValueError):
        rs.scale_factor(4)
  
    def test_block_name(self):
      self.assertEqual(networks.block_name(10), 'progressive_gan_block_10')
  
    def test_min_total_num_images(self):
      self.assertEqual(networks.min_total_num_images(7, 8, 4), 52)
  
    def test_compute_progress(self):
      if tf.executing_eagerly():
        progress_output = []
        for current_image_id in [0, 3, 6, 7, 8, 10, 15, 29, 100]:
          progress = networks.compute_progress(
              current_image_id,
              stable_stage_num_images=7,
              transition_stage_num_images=8,
              num_blocks=2)
          with self.cached_session(use_gpu=True) as sess:
            progress_output.append(sess.run(progress))
      else:
        current_image_id_ph = tf.compat.v1.placeholder(tf.int32, [])
        progress = networks.compute_progress(
            current_image_id_ph,
            stable_stage_num_images=7,
            transition_stage_num_images=8,
            num_blocks=2)
        with self.cached_session(use_gpu=True) as sess:
          progress_output = [
              sess.run(progress, feed_dict={current_image_id_ph: cur_image_id})
              for cur_image_id in [0, 3, 6, 7, 8, 10, 15, 29, 100]
          ]
  
      self.assertArrayNear(progress_output,
                           [0.0, 0.0, 0.0, 0.0, 0.125, 0.375, 1.0, 1.0, 1.0],
                           1.0e-6)
  
    def test_generator_alpha(self):
      with self.cached_session(use_gpu=True) as sess:
        alpha_fixed_block_id = [
            sess.run(
                networks._generator_alpha(2, tf.constant(progress, tf.float32)))
            for progress in [0, 0.2, 1, 1.2, 2, 2.2, 3]
        ]
        alpha_fixed_progress = [
            sess.run(
                networks._generator_alpha(block_id, tf.constant(1.2, tf.float32)))
            for block_id in range(1, 5)
        ]
  
      self.assertArrayNear(alpha_fixed_block_id, [0, 0.2, 1, 0.8, 0, 0, 0],
                           1.0e-6)
      self.assertArrayNear(alpha_fixed_progress, [0, 0.8, 0.2, 0], 1.0e-6)
  
    def test_discriminator_alpha(self):
      with self.cached_session(use_gpu=True) as sess:
        alpha_fixed_block_id = [sess.run(networks._discriminator_alpha(
            2, tf.constant(progress, tf.float32))) for progress in
                                [0, 0.2, 1, 1.2, 2, 2.2, 3]]
        alpha_fixed_progress = [sess.run(networks._discriminator_alpha(
            block_id, tf.constant(1.2, tf.float32))) for block_id in range(1, 5)]
  
      self.assertArrayNear(alpha_fixed_block_id, [1, 1, 1, 0.8, 0, 0, 0], 1.0e-6)
      self.assertArrayNear(alpha_fixed_progress, [0, 0.8, 1, 1], 1.0e-6)
  
    def test_blend_images_in_stable_stage(self):
      x_np = np.random.normal(size=[2, 8, 8, 3])
      x = tf.constant(x_np, tf.float32)
      x_blend = networks.blend_images(
          x,
          progress=tf.constant(0.0),
          resolution_schedule=networks.ResolutionSchedule(
              scale_base=2, num_resolutions=2),
          num_blocks=2)
      with self.cached_session(use_gpu=True) as sess:
        x_blend_np = sess.run(x_blend)
        x_blend_expected_np = sess.run(layers.upscale(layers.downscale(x, 2), 2))
      self.assertNDArrayNear(x_blend_np, x_blend_expected_np, 1.0e-6)
  
    def test_blend_images_in_transition_stage(self):
      x_np = np.random.normal(size=[2, 8, 8, 3])
      x = tf.constant(x_np, tf.float32)
      x_blend = networks.blend_images(
          x,
          tf.constant(0.2),
          resolution_schedule=networks.ResolutionSchedule(
              scale_base=2, num_resolutions=2),
          num_blocks=2)
      with self.cached_session(use_gpu=True) as sess:
        x_blend_np = sess.run(x_blend)
        x_blend_expected_np = 0.8 * sess.run(
            layers.upscale(layers.downscale(x, 2), 2)) + 0.2 * x_np
      self.assertNDArrayNear(x_blend_np, x_blend_expected_np, 1.0e-6)
  
    def test_num_filters(self):
      self.assertEqual(networks.num_filters(1, 4096, 1, 256), 256)
      self.assertEqual(networks.num_filters(5, 4096, 1, 256), 128)
  
    def test_generator_grad_norm_progress(self):
      if tf.executing_eagerly():
        # tf.placeholder() is not compatible with eager execution.
        return
      stable_stage_num_images = 2
      transition_stage_num_images = 3
  
      current_image_id_ph = tf.compat.v1.placeholder(tf.int32, [])
      progress = networks.compute_progress(
          current_image_id_ph,
          stable_stage_num_images,
          transition_stage_num_images,
          num_blocks=3)
      z = tf.random.normal([2, 10], dtype=tf.float32)
      x, _ = networks.generator(
          z, progress, _num_filters_stub,
          networks.ResolutionSchedule(
              start_resolutions=(4, 4), scale_base=2, num_resolutions=3))
      fake_loss = tf.reduce_sum(input_tensor=tf.square(x))
      grad_norms = [
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_1/.*')),
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_2/.*')),
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_3/.*'))
      ]
  
      grad_norms_output = None
      with self.cached_session(use_gpu=True) as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        x1_np = sess.run(x, feed_dict={current_image_id_ph: 0.12})
        x2_np = sess.run(x, feed_dict={current_image_id_ph: 1.8})
        grad_norms_output = np.array([
            sess.run(grad_norms, feed_dict={current_image_id_ph: i})
            for i in range(15)  # total num of images
        ])
  
      self.assertEqual((2, 16, 16, 3), x1_np.shape)
      self.assertEqual((2, 16, 16, 3), x2_np.shape)
      # The gradient of block_1 is always on.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 0] > 0), 0,
          'gradient norms {} for block 1 is not always on'.format(
              grad_norms_output[:, 0]))
      # The gradient of block_2 is on after 1 stable stage.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 1] > 0), 3,
          'gradient norms {} for block 2 is not on at step 3'.format(
              grad_norms_output[:, 1]))
      # The gradient of block_3 is on after 2 stable stage + 1 transition stage.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 2] > 0), 8,
          'gradient norms {} for block 3 is not on at step 8'.format(
              grad_norms_output[:, 2]))
  
    def test_discriminator_grad_norm_progress(self):
      if tf.executing_eagerly():
        # tf.placeholder() is not compatible with eager execution.
        return
      stable_stage_num_images = 2
      transition_stage_num_images = 3
  
      current_image_id_ph = tf.compat.v1.placeholder(tf.int32, [])
      progress = networks.compute_progress(
          current_image_id_ph,
          stable_stage_num_images,
          transition_stage_num_images,
          num_blocks=3)
      x = tf.random.normal([2, 16, 16, 3])
      logits, _ = networks.discriminator(
          x, progress, _num_filters_stub,
          networks.ResolutionSchedule(
              start_resolutions=(4, 4), scale_base=2, num_resolutions=3))
      fake_loss = tf.reduce_sum(input_tensor=tf.square(logits))
      grad_norms = [
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_1/.*')),
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_2/.*')),
          _get_grad_norm(
              fake_loss,
              tf.compat.v1.trainable_variables('.*/progressive_gan_block_3/.*'))
      ]
  
      grad_norms_output = None
      with self.cached_session(use_gpu=True) as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        grad_norms_output = np.array([
            sess.run(grad_norms, feed_dict={current_image_id_ph: i})
            for i in range(15)  # total num of images
        ])
  
      # The gradient of block_1 is always on.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 0] > 0), 0,
          'gradient norms {} for block 1 is not always on'.format(
              grad_norms_output[:, 0]))
      # The gradient of block_2 is on after 1 stable stage.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 1] > 0), 3,
          'gradient norms {} for block 2 is not on at step 3'.format(
              grad_norms_output[:, 1]))
      # The gradient of block_3 is on after 2 stable stage + 1 transition stage.
      self.assertEqual(
          np.argmax(grad_norms_output[:, 2] > 0), 8,
          'gradient norms {} for block 3 is not on at step 8'.format(
              grad_norms_output[:, 2]))
  
  
  if __name__ == '__main__':
    tf.test.main()
  
apis: tensorflow.compat.v1.placeholder numpy.argmax tensorflow.compat.v1.trainable_variables tensorflow.gradients tensorflow.compat.v1.global_variables_initializer tensorflow.square tensorflow.random.normal numpy.random.normal tensorflow.executing_eagerly tensorflow.constant tensorflow.test.main