Body: <p>I'm trying to implement and train a Progressive GAN for TensorFlow (to generate faces like the CelebA dataset), following along with the original paper (<a href="https://arxiv.org/pdf/1710.10196.pdf" rel="nofollow noreferrer">https://arxiv.org/pdf/1710.10196.pdf</a>), as well as occasionally referring to this video guide for implementing it in PyTorch <a href="https://www.youtube.com/watch?v=nkQHASviYac" rel="nofollow noreferrer">https://www.youtube.com/watch?v=nkQHASviYac</a></p>
  <p>However, after entire days of trying, I can't manage to get a stabilized training until the end. For now I'm going only until 64x64 images, but training often destabilizes as soon as I start fading in a layer. Weird thing is that the training restabilizes one epoch later if I just add a layer abruptly. This at first led me to believe that there must be some operation that is failing to get gradients properly because of the way I am calculating the fadein, and this was supported by some asserts being caught showing that weights are being updated to NaN. However, after a correction in my call() functions for both networks, those errors disappeared (unless I try to add spectral normalization, for whatever reason). Now, training just seems to be just unstable and the issues I'm facing might be exploding gradients, and I'm honestly just really confused.</p>
  <p>I'm wondering if someone experienced can see some obvious errors in how I'm implementing this progressive GAN which is leading to problems with calculating gradients, paying key attention to how I'm fading in new layers and adding them with the previous ones, or if there's general architecture problems leading to instability.</p>
  <p>Any help or tips would be appreciated, thank you so much in advance.</p>
  <pre><code>    import time
      import tensorflow as tf
      import os
      import matplotlib.pyplot as plt
      from keras import Input
      from keras.layers import Dense, BatchNormalization, LeakyReLU, Reshape, Conv2DTranspose, Conv2D, Dropout, ReLU, Flatten, \
          UpSampling2D, AveragePooling2D
      from tensorflow_addons.layers import SpectralNormalization
      import tensorflow_datasets as tfds
      import numpy as np
      
      # tf.config.run_functions_eagerly(True)
      
      BATCH_SIZE = 16
      EPOCHS = 5000
      IMAGE_DIMS = (128, 128)
      GENERATOR_LR = 2.5e-4
      DISCRIMINATOR_LR = 2.5e-4
      STARTING_STEP = 0
      FADEIN_AMNT = 0.01
      FADEIN_EPOCHS = 10  # How often we FadeIn
      
      # Load the CelebA dataset
      celebA, info = tfds.load(&quot;celeb_a&quot;, with_info=True)
      train_data = celebA[&quot;train&quot;]
      
      num_samples = info.splits['train'].num_examples
      
      
      def load_image(datapoint):
          image = datapoint['image']
          image = tf.image.resize(image, IMAGE_DIMS)  # Resize the image to a standard size
          image = (tf.cast(image, tf.float32) - 127.5) / 127.5  # Normalize the pixel values to [-1, 1]
          return image
      
      
      image_dataset = train_data.map(load_image)
      image_dataset = image_dataset.batch(BATCH_SIZE, drop_remainder=True).take(50)
      
      # Create the seed for image generation
      NOISE_DIM = 256
      noise = tf.random.normal([1, NOISE_DIM])
      
      
      def equalized_learning_rate(shape, dtype=tf.float32, partition_info=None):
          &quot;&quot;&quot;This initializes the layer weights and applies equalized learning rate&quot;&quot;&quot;
          if len(shape) == 2:  # It is a dense layer
              fan_in = shape[0]
          elif len(shape) &gt; 2:  # It is a conv layer
              fan_in = np.prod(shape[:-1])  # product of dimensions
          return tf.random.normal(shape=shape, mean=0., stddev=np.sqrt(2 / fan_in), dtype=dtype)
      
      
      class EqualizedDense(tf.keras.layers.Dense):
          def __init__(self, units, **kwargs):
              super().__init__(units, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class EqualizedConv2D(tf.keras.layers.Conv2D):
          def __init__(self, filters, kernel_size, **kwargs):
              super().__init__(filters, kernel_size, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class EqualizedConv2DTranspose(tf.keras.layers.Conv2DTranspose):
          def __init__(self, filters, kernel_size, **kwargs):
              super().__init__(filters, kernel_size, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class PixelNormalization(tf.keras.layers.Layer):
          &quot;&quot;&quot;Divides pixels by their squared means along a batch&quot;&quot;&quot;
      
          def __init__(self, epsilon=1e-8):
              super(PixelNormalization, self).__init__()
              self.epsilon = epsilon
      
          def call(self, inputs, *args, **kwargs):
              return inputs / tf.math.sqrt(tf.reduce_mean(tf.square(inputs), axis=-1, keepdims=True) + self.epsilon)
      
      
      class MinibatchStddev(tf.keras.layers.Layer):
          def call(self, x, **kwargs):
              stddev = tf.math.reduce_std(x, axis=0, keepdims=True)
              stddev = tf.reduce_mean(stddev, axis=[1, 2, 3], keepdims=True)
              stddev = tf.tile(stddev, multiples=[tf.shape(x)[0], x.shape[1], x.shape[2], 1])
              return tf.concat([x, stddev], axis=-1)
      
      
      def call_block(x, block):
          for layer in block:
              x = layer(x)
          return x
      
      
      class ProgressiveGenerator(tf.keras.Model):
          def __init__(self):
              super(ProgressiveGenerator, self).__init__()
      
              initial_filters = 256
              self.filters = [256, 128, 64, 32]  # excludes initial block
      
              self.initial_block = [
                  Dense(4 * 4 * initial_filters),
                  Reshape((4, 4, initial_filters)),
                  PixelNormalization(),
                  Conv2DTranspose(initial_filters, 4, padding='same'),
                  LeakyReLU(0.2),
                  EqualizedConv2D(initial_filters, 3, padding='same'),
                  LeakyReLU(0.2),
                  PixelNormalization()
              ]
      
              self.progressive_blocks = []
              self.to_rgb_blocks = [EqualizedConv2D(3, 1, padding=&quot;same&quot;, activation='tanh')]
      
              for filters in self.filters:
                  self.progressive_blocks.append(
                      [
                          UpSampling2D(),
                          EqualizedConv2D(filters, 3, padding='same'),
                          EqualizedConv2D(filters, 3, padding='same'),
                          LeakyReLU(0.2),
                          PixelNormalization()
                      ]
                  )
                  self.to_rgb_blocks.append(EqualizedConv2D(3, 1, padding=&quot;same&quot;, activation='tanh'))
      
          def call(self, inputs, current_step=0, alpha=0.2):
              x = inputs
              x = call_block(x, self.initial_block)
              if current_step == 0:
                  return self.to_rgb_blocks[0](x)
              else:
                  prev = None
                  for step in range(current_step):
                      if step == current_step - 1:
                          prev = self.to_rgb_blocks[step](UpSampling2D()(x))  # Upsample before the to_rgb operation
                      x = call_block(x, self.progressive_blocks[step])
                  x = self.to_rgb_blocks[current_step](x)
                  x = tf.keras.layers.add([alpha * x, (1 - alpha) * prev])
                  assert not tf.math.is_nan(tf.reduce_sum(x)), &quot;Nan detected in x&quot;
                  return x
      
      
      class ProgressiveDiscriminator(tf.keras.Model):
          def __init__(self):
              super(ProgressiveDiscriminator, self).__init__()
      
              self.filters = [32, 64, 128, 256]
              initial_filters = 256
      
              self.initial_from_rgb = SpectralNormalization(EqualizedConv2D(initial_filters, 1, padding='same'))
              self.from_rgb = []
      
              self.progressive_blocks = []
      
              for filters in self.filters:
                  self.progressive_blocks.append([
                      AveragePooling2D(),
                      SpectralNormalization(EqualizedConv2D(filters, 3, padding='same')),
                      SpectralNormalization(EqualizedConv2D(filters, 3, padding='same')),
                      LeakyReLU(0.2),
                  ])
                  self.from_rgb.append(SpectralNormalization(EqualizedConv2D(filters // 2, 1, padding='same')))
      
              self.output_block = [
                  MinibatchStddev(),
                  LeakyReLU(0.2),
                  SpectralNormalization(EqualizedConv2D(initial_filters, 3, padding='same')),
                  SpectralNormalization(EqualizedConv2D(1, 1, padding='same')),
                  Reshape((16,)),
                  Dense(1)
              ]
      
          def call(self, inputs, current_step=0, alpha=0.2):
              x = inputs
              if current_step == 0:
                  x = self.initial_from_rgb(x)
                  return call_block(x, self.output_block)
              else:
                  downscaled = AveragePooling2D()(x)
                  if current_step == 1:
                      downscaled = self.initial_from_rgb(downscaled)
                  else:
                      downscaled = self.from_rgb[len(self.progressive_blocks) - current_step + 1](downscaled)
      
                  x = self.from_rgb[len(self.progressive_blocks) - current_step](x)
                  for step in reversed(range(current_step)):
                      x = call_block(x, self.progressive_blocks[len(self.progressive_blocks) - step - 1])
                      if step == current_step - 1:
                          x = tf.keras.layers.add([alpha * x, (1 - alpha) * downscaled])
                  x = call_block(x, self.output_block)
                  return x
      
      
      generator = ProgressiveGenerator()
      generator.build((None, NOISE_DIM))
      
      discriminator = ProgressiveDiscriminator()
      discriminator.build((None, 4, 4, 3))
      
      generator.summary()
      discriminator.summary()
      
      generator_optimizer = tf.keras.optimizers.legacy.Adam(GENERATOR_LR, beta_1=0., beta_2=0.99)
      discriminator_optimizer = tf.keras.optimizers.legacy.Adam(DISCRIMINATOR_LR, beta_1=0., beta_2=0.99)
      
      generator.compile(optimizer=generator_optimizer)
      discriminator.compile(optimizer=discriminator_optimizer)
      
      
      def grad_penalty(real_images, fake_images, current_step):
          batch_size = tf.shape(real_images)[0]
          epsilon = tf.random.uniform((batch_size, 1, 1, 1), 0.0, 1.0)
      
          real_images_resized = tf.image.resize(real_images, [fake_images.shape[1], fake_images.shape[2]])
          interpolated_images = epsilon * real_images_resized + ((1 - epsilon) * fake_images)
      
          with tf.GradientTape() as tape:
              tape.watch(interpolated_images)
              critic_interpolated = discriminator(interpolated_images, current_step=current_step)
      
          gradients = tape.gradient(critic_interpolated, [interpolated_images])[0]
          gradients_norm = tf.sqrt(tf.reduce_sum(tf.square(gradients), axis=[1, 2, 3]))
          gradient_penalty = tf.reduce_mean((gradients_norm - 1.) ** 2)
      
          return gradient_penalty
      
      
      def train_step(images, current_step, alpha):
          batch_size = tf.shape(images)[0]
          current_resolution = 4 * (2 ** current_step)
          images_resized = tf.image.resize(images, [current_resolution, current_resolution])
      
          n_critic = 1
          for _ in range(n_critic):
              with tf.GradientTape() as disc_tape:
                  noise = tf.random.normal([batch_size, NOISE_DIM])
                  generated_images = generator(noise, current_step=current_step, alpha=alpha, training=True)
      
                  assert not tf.math.is_nan(tf.reduce_sum(generated_images)), &quot;Nan detected in generated_images&quot;
      
                  real_predictions = discriminator(images_resized, current_step=current_step, alpha=alpha, training=True)
                  fake_predictions = discriminator(generated_images, current_step=current_step, alpha=alpha, training=True)
      
                  assert not tf.math.is_nan(tf.reduce_sum(real_predictions)), &quot;Nan detected in real_predictions&quot;
                  assert not tf.math.is_nan(tf.reduce_sum(fake_predictions)), &quot;Nan detected in fake_predictions&quot;
      
                  gradient_penalty = grad_penalty(images_resized, generated_images, current_step)
                  drift_term = 0.001 * tf.reduce_mean(tf.square(real_predictions))
                  disc_loss = tf.reduce_mean(fake_predictions - real_predictions + 10. * gradient_penalty + drift_term)
      
                  disc_grad = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
                  disc_grad = [(g, v) for g, v in zip(disc_grad, discriminator.trainable_variables) if g is not None]
                  discriminator_optimizer.apply_gradients(disc_grad)
      
          noise = tf.random.normal([batch_size, NOISE_DIM])
          with tf.GradientTape() as gen_tape:
              generated_images = generator(noise, current_step=current_step, training=True)
              fake_output = discriminator(generated_images, current_step=current_step, training=True)
              gen_loss = -tf.reduce_mean(fake_output)
      
          gen_grad = gen_tape.gradient(gen_loss, generator.trainable_variables)
          gen_grad = [(g, v) for g, v in zip(gen_grad, generator.trainable_variables) if g is not None]
          generator_optimizer.apply_gradients(gen_grad)
          return gen_loss, disc_loss
      
      
      def train(dataset, epochs):
          seed = tf.random.normal([16, NOISE_DIM])
      
          alpha = 0.0
          current_step = STARTING_STEP
      
          for epoch in range(epochs):
              start = time.time()
              gen_losses = []
              disc_losses = []
      
              if alpha &gt;= 0.99:
                  current_step = min(current_step + 1, 4)  # increased after each fadein
                  alpha = 0.0
      
              for image_batch in dataset:
                  gen_loss, disc_loss = train_step(image_batch, current_step, alpha)
                  gen_losses.append(gen_loss)
                  disc_losses.append(disc_loss)
      
              print(&quot;alpha: &quot;, alpha)
              generate_and_save_images(generator, current_step, epoch + 1, seed)
      
              if (epoch + 1) % 15 == 0:
                  print(&quot;Saving models...&quot;)
                  checkpoint.save(file_prefix=checkpoint_prefix)
              if (epoch + 1) % FADEIN_EPOCHS == 0:
                  alpha += FADEIN_AMNT
      
              avg_gen_loss = sum(gen_losses) / len(gen_losses)
              avg_disc_loss = sum(disc_losses) / len(disc_losses)
              print(f'Epoch {epoch + 1} - Generator Loss: {avg_gen_loss}, Discriminator Loss: {avg_disc_loss}')
              print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))
      
      
      def generate_and_save_images(model, current_step, epoch, test_input):
          predictions = model(test_input, current_step=current_step, training=False)
      
          for i in range(predictions.shape[0]):
              plt.subplot(4, 4, i + 1)
              plt.imshow((predictions[i] * 0.5 + 0.5))
              plt.axis('off')
          print('Saving image at resolution step', current_step)
          plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
      
      
      checkpoint_dir = './training_checkpoints'
      checkpoint_prefix = os.path.join(checkpoint_dir, &quot;ckpt&quot;)
      checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                       discriminator_optimizer=discriminator_optimizer,
                                       generator=generator,
                                       discriminator=discriminator)
      
      train(image_dataset, EPOCHS)<span class="math-container">```</span>
  </code></pre>
  
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-04-23T21:27:47.530Z
Comments.Text: This isn’t an Answer, but a request for help as well. I have been facing the same issue pretty much. If I fade new layers in slowly over many epochs, the losses start to go crazy for a for epochs. It does recover for lower resolutions but doesn’t for higher ones (>32). And if I fade in quickly within 1 epoch, there is almost no change in the loss but it then just gets stuck at the min values in a very suspicious way. The output images seem to be fine for lower resolutions, but the loss values are just very very small, and do not seem to increase or decrease. Infact after a higher layer is adde
Comments.UserDisplayName: Ashish Kumar Singh
Comments.ContentLicense: CC BY-SA 4.0
Comments.CreationDate: 2024-05-01T19:07:06.077Z
Comments.Text: @AshishKumarSingh I never figured it out, I ended up switching to PyTorch instead and it yielded far better results. TensorFlow can be weird at times.
Comments.UserDisplayName: alexdea
Tags: <generative-adversarial-networks>
Title: Unstable Progressive GAN Training with Exploding Gradients when Fading in New Layers
UserDisplayName: alexdea
ThreadText: # Unstable Progressive GAN Training with Exploding Gradients when Fading in New Layers
  
  (Asked by: username_0 on 2023-06-22)
  
  I'm trying to implement and train a Progressive GAN for TensorFlow (to generate faces like the CelebA dataset), following along with the original paper (<https://arxiv.org/pdf/1710.10196.pdf>), as well as occasionally referring to this video guide for implementing it in PyTorch <https://www.youtube.com/watch?v=nkQHASviYac>
  
  However, after entire days of trying, I can't manage to get a stabilized training until the end. For now I'm going only until 64x64 images, but training often destabilizes as soon as I start fading in a layer. Weird thing is that the training restabilizes one epoch later if I just add a layer abruptly. This at first led me to believe that there must be some operation that is failing to get gradients properly because of the way I am calculating the fadein, and this was supported by some asserts being caught showing that weights are being updated to NaN. However, after a correction in my call() functions for both networks, those errors disappeared (unless I try to add spectral normalization, for whatever reason). Now, training just seems to be just unstable and the issues I'm facing might be exploding gradients, and I'm honestly just really confused.
  
  I'm wondering if someone experienced can see some obvious errors in how I'm implementing this progressive GAN which is leading to problems with calculating gradients, paying key attention to how I'm fading in new layers and adding them with the previous ones, or if there's general architecture problems leading to instability.
  
  Any help or tips would be appreciated, thank you so much in advance.
  
  ```
      import time
      import tensorflow as tf
      import os
      import matplotlib.pyplot as plt
      from keras import Input
      from keras.layers import Dense, BatchNormalization, LeakyReLU, Reshape, Conv2DTranspose, Conv2D, Dropout, ReLU, Flatten, \
          UpSampling2D, AveragePooling2D
      from tensorflow_addons.layers import SpectralNormalization
      import tensorflow_datasets as tfds
      import numpy as np
      
      # tf.config.run_functions_eagerly(True)
      
      BATCH_SIZE = 16
      EPOCHS = 5000
      IMAGE_DIMS = (128, 128)
      GENERATOR_LR = 2.5e-4
      DISCRIMINATOR_LR = 2.5e-4
      STARTING_STEP = 0
      FADEIN_AMNT = 0.01
      FADEIN_EPOCHS = 10  # How often we FadeIn
      
      # Load the CelebA dataset
      celebA, info = tfds.load("celeb_a", with_info=True)
      train_data = celebA["train"]
      
      num_samples = info.splits['train'].num_examples
      
      
      def load_image(datapoint):
          image = datapoint['image']
          image = tf.image.resize(image, IMAGE_DIMS)  # Resize the image to a standard size
          image = (tf.cast(image, tf.float32) - 127.5) / 127.5  # Normalize the pixel values to [-1, 1]
          return image
      
      
      image_dataset = train_data.map(load_image)
      image_dataset = image_dataset.batch(BATCH_SIZE, drop_remainder=True).take(50)
      
      # Create the seed for image generation
      NOISE_DIM = 256
      noise = tf.random.normal([1, NOISE_DIM])
      
      
      def equalized_learning_rate(shape, dtype=tf.float32, partition_info=None):
          """This initializes the layer weights and applies equalized learning rate"""
          if len(shape) == 2:  # It is a dense layer
              fan_in = shape[0]
          elif len(shape) > 2:  # It is a conv layer
              fan_in = np.prod(shape[:-1])  # product of dimensions
          return tf.random.normal(shape=shape, mean=0., stddev=np.sqrt(2 / fan_in), dtype=dtype)
      
      
      class EqualizedDense(tf.keras.layers.Dense):
          def __init__(self, units, **kwargs):
              super().__init__(units, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class EqualizedConv2D(tf.keras.layers.Conv2D):
          def __init__(self, filters, kernel_size, **kwargs):
              super().__init__(filters, kernel_size, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class EqualizedConv2DTranspose(tf.keras.layers.Conv2DTranspose):
          def __init__(self, filters, kernel_size, **kwargs):
              super().__init__(filters, kernel_size, kernel_initializer=equalized_learning_rate, **kwargs)
      
      
      class PixelNormalization(tf.keras.layers.Layer):
          """Divides pixels by their squared means along a batch"""
      
          def __init__(self, epsilon=1e-8):
              super(PixelNormalization, self).__init__()
              self.epsilon = epsilon
      
          def call(self, inputs, *args, **kwargs):
              return inputs / tf.math.sqrt(tf.reduce_mean(tf.square(inputs), axis=-1, keepdims=True) + self.epsilon)
      
      
      class MinibatchStddev(tf.keras.layers.Layer):
          def call(self, x, **kwargs):
              stddev = tf.math.reduce_std(x, axis=0, keepdims=True)
              stddev = tf.reduce_mean(stddev, axis=[1, 2, 3], keepdims=True)
              stddev = tf.tile(stddev, multiples=[tf.shape(x)[0], x.shape[1], x.shape[2], 1])
              return tf.concat([x, stddev], axis=-1)
      
      
      def call_block(x, block):
          for layer in block:
              x = layer(x)
          return x
      
      
      class ProgressiveGenerator(tf.keras.Model):
          def __init__(self):
              super(ProgressiveGenerator, self).__init__()
      
              initial_filters = 256
              self.filters = [256, 128, 64, 32]  # excludes initial block
      
              self.initial_block = [
                  Dense(4 * 4 * initial_filters),
                  Reshape((4, 4, initial_filters)),
                  PixelNormalization(),
                  Conv2DTranspose(initial_filters, 4, padding='same'),
                  LeakyReLU(0.2),
                  EqualizedConv2D(initial_filters, 3, padding='same'),
                  LeakyReLU(0.2),
                  PixelNormalization()
              ]
      
              self.progressive_blocks = []
              self.to_rgb_blocks = [EqualizedConv2D(3, 1, padding="same", activation='tanh')]
      
              for filters in self.filters:
                  self.progressive_blocks.append(
                      [
                          UpSampling2D(),
                          EqualizedConv2D(filters, 3, padding='same'),
                          EqualizedConv2D(filters, 3, padding='same'),
                          LeakyReLU(0.2),
                          PixelNormalization()
                      ]
                  )
                  self.to_rgb_blocks.append(EqualizedConv2D(3, 1, padding="same", activation='tanh'))
      
          def call(self, inputs, current_step=0, alpha=0.2):
              x = inputs
              x = call_block(x, self.initial_block)
              if current_step == 0:
                  return self.to_rgb_blocks[0](x)
              else:
                  prev = None
                  for step in range(current_step):
                      if step == current_step - 1:
                          prev = self.to_rgb_blocks[step](UpSampling2D()(x))  # Upsample before the to_rgb operation
                      x = call_block(x, self.progressive_blocks[step])
                  x = self.to_rgb_blocks[current_step](x)
                  x = tf.keras.layers.add([alpha * x, (1 - alpha) * prev])
                  assert not tf.math.is_nan(tf.reduce_sum(x)), "Nan detected in x"
                  return x
      
      
      class ProgressiveDiscriminator(tf.keras.Model):
          def __init__(self):
              super(ProgressiveDiscriminator, self).__init__()
      
              self.filters = [32, 64, 128, 256]
              initial_filters = 256
      
              self.initial_from_rgb = SpectralNormalization(EqualizedConv2D(initial_filters, 1, padding='same'))
              self.from_rgb = []
      
              self.progressive_blocks = []
      
              for filters in self.filters:
                  self.progressive_blocks.append([
                      AveragePooling2D(),
                      SpectralNormalization(EqualizedConv2D(filters, 3, padding='same')),
                      SpectralNormalization(EqualizedConv2D(filters, 3, padding='same')),
                      LeakyReLU(0.2),
                  ])
                  self.from_rgb.append(SpectralNormalization(EqualizedConv2D(filters // 2, 1, padding='same')))
      
              self.output_block = [
                  MinibatchStddev(),
                  LeakyReLU(0.2),
                  SpectralNormalization(EqualizedConv2D(initial_filters, 3, padding='same')),
                  SpectralNormalization(EqualizedConv2D(1, 1, padding='same')),
                  Reshape((16,)),
                  Dense(1)
              ]
      
          def call(self, inputs, current_step=0, alpha=0.2):
              x = inputs
              if current_step == 0:
                  x = self.initial_from_rgb(x)
                  return call_block(x, self.output_block)
              else:
                  downscaled = AveragePooling2D()(x)
                  if current_step == 1:
                      downscaled = self.initial_from_rgb(downscaled)
                  else:
                      downscaled = self.from_rgb[len(self.progressive_blocks) - current_step + 1](downscaled)
      
                  x = self.from_rgb[len(self.progressive_blocks) - current_step](x)
                  for step in reversed(range(current_step)):
                      x = call_block(x, self.progressive_blocks[len(self.progressive_blocks) - step - 1])
                      if step == current_step - 1:
                          x = tf.keras.layers.add([alpha * x, (1 - alpha) * downscaled])
                  x = call_block(x, self.output_block)
                  return x
      
      
      generator = ProgressiveGenerator()
      generator.build((None, NOISE_DIM))
      
      discriminator = ProgressiveDiscriminator()
      discriminator.build((None, 4, 4, 3))
      
      generator.summary()
      discriminator.summary()
      
      generator_optimizer = tf.keras.optimizers.legacy.Adam(GENERATOR_LR, beta_1=0., beta_2=0.99)
      discriminator_optimizer = tf.keras.optimizers.legacy.Adam(DISCRIMINATOR_LR, beta_1=0., beta_2=0.99)
      
      generator.compile(optimizer=generator_optimizer)
      discriminator.compile(optimizer=discriminator_optimizer)
      
      
      def grad_penalty(real_images, fake_images, current_step):
          batch_size = tf.shape(real_images)[0]
          epsilon = tf.random.uniform((batch_size, 1, 1, 1), 0.0, 1.0)
      
          real_images_resized = tf.image.resize(real_images, [fake_images.shape[1], fake_images.shape[2]])
          interpolated_images = epsilon * real_images_resized + ((1 - epsilon) * fake_images)
      
          with tf.GradientTape() as tape:
              tape.watch(interpolated_images)
              critic_interpolated = discriminator(interpolated_images, current_step=current_step)
      
          gradients = tape.gradient(critic_interpolated, [interpolated_images])[0]
          gradients_norm = tf.sqrt(tf.reduce_sum(tf.square(gradients), axis=[1, 2, 3]))
          gradient_penalty = tf.reduce_mean((gradients_norm - 1.) ** 2)
      
          return gradient_penalty
      
      
      def train_step(images, current_step, alpha):
          batch_size = tf.shape(images)[0]
          current_resolution = 4 * (2 ** current_step)
          images_resized = tf.image.resize(images, [current_resolution, current_resolution])
      
          n_critic = 1
          for _ in range(n_critic):
              with tf.GradientTape() as disc_tape:
                  noise = tf.random.normal([batch_size, NOISE_DIM])
                  generated_images = generator(noise, current_step=current_step, alpha=alpha, training=True)
      
                  assert not tf.math.is_nan(tf.reduce_sum(generated_images)), "Nan detected in generated_images"
      
                  real_predictions = discriminator(images_resized, current_step=current_step, alpha=alpha, training=True)
                  fake_predictions = discriminator(generated_images, current_step=current_step, alpha=alpha, training=True)
      
                  assert not tf.math.is_nan(tf.reduce_sum(real_predictions)), "Nan detected in real_predictions"
                  assert not tf.math.is_nan(tf.reduce_sum(fake_predictions)), "Nan detected in fake_predictions"
      
                  gradient_penalty = grad_penalty(images_resized, generated_images, current_step)
                  drift_term = 0.001 * tf.reduce_mean(tf.square(real_predictions))
                  disc_loss = tf.reduce_mean(fake_predictions - real_predictions + 10. * gradient_penalty + drift_term)
      
                  disc_grad = disc_tape.gradient(disc_loss, discriminator.trainable_variables)
                  disc_grad = [(g, v) for g, v in zip(disc_grad, discriminator.trainable_variables) if g is not None]
                  discriminator_optimizer.apply_gradients(disc_grad)
      
          noise = tf.random.normal([batch_size, NOISE_DIM])
          with tf.GradientTape() as gen_tape:
              generated_images = generator(noise, current_step=current_step, training=True)
              fake_output = discriminator(generated_images, current_step=current_step, training=True)
              gen_loss = -tf.reduce_mean(fake_output)
      
          gen_grad = gen_tape.gradient(gen_loss, generator.trainable_variables)
          gen_grad = [(g, v) for g, v in zip(gen_grad, generator.trainable_variables) if g is not None]
          generator_optimizer.apply_gradients(gen_grad)
          return gen_loss, disc_loss
      
      
      def train(dataset, epochs):
          seed = tf.random.normal([16, NOISE_DIM])
      
          alpha = 0.0
          current_step = STARTING_STEP
      
          for epoch in range(epochs):
              start = time.time()
              gen_losses = []
              disc_losses = []
      
              if alpha >= 0.99:
                  current_step = min(current_step + 1, 4)  # increased after each fadein
                  alpha = 0.0
      
              for image_batch in dataset:
                  gen_loss, disc_loss = train_step(image_batch, current_step, alpha)
                  gen_losses.append(gen_loss)
                  disc_losses.append(disc_loss)
      
              print("alpha: ", alpha)
              generate_and_save_images(generator, current_step, epoch + 1, seed)
      
              if (epoch + 1) % 15 == 0:
                  print("Saving models...")
                  checkpoint.save(file_prefix=checkpoint_prefix)
              if (epoch + 1) % FADEIN_EPOCHS == 0:
                  alpha += FADEIN_AMNT
      
              avg_gen_loss = sum(gen_losses) / len(gen_losses)
              avg_disc_loss = sum(disc_losses) / len(disc_losses)
              print(f'Epoch {epoch + 1} - Generator Loss: {avg_gen_loss}, Discriminator Loss: {avg_disc_loss}')
              print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))
      
      
      def generate_and_save_images(model, current_step, epoch, test_input):
          predictions = model(test_input, current_step=current_step, training=False)
      
          for i in range(predictions.shape[0]):
              plt.subplot(4, 4, i + 1)
              plt.imshow((predictions[i] * 0.5 + 0.5))
              plt.axis('off')
          print('Saving image at resolution step', current_step)
          plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))
      
      
      checkpoint_dir = './training_checkpoints'
      checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
      checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                       discriminator_optimizer=discriminator_optimizer,
                                       generator=generator,
                                       discriminator=discriminator)
      
      train(image_dataset, EPOCHS)```
  
  ```
  
  **Tags:** generative-adversarial-networks
  
  > **Comments:**
  > **username_1**: This isn’t an Answer, but a request for help as well. I have been facing the same issue pretty much. If I fade new layers in slowly over many epochs, the losses start to go crazy for a for epochs. It does recover for lower resolutions but doesn’t for higher ones (>32). And if I fade in quickly within 1 epoch, there is almost no change in the loss but it then just gets stuck at the min values in a very suspicious way. The output images seem to be fine for lower resolutions, but the loss values are just very very small, and do not seem to increase or decrease. Infact after a higher layer is adde  
  > **username_0**: @username_1 I never figured it out, I ended up switching to PyTorch instead and it yielded far better results. TensorFlow can be weird at times.  
  
  ---
  
  ## 0 answers
  
  