<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>

CC BY-SA 4.0
2024-04-23T21:27:47.530Z
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
Ashish Kumar Singh
CC BY-SA 4.0
2024-05-01T19:07:06.077Z
@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.
alexdea
<generative-adversarial-networks>
Unstable Progressive GAN Training with Exploding Gradients when Fading in New Layers
alexdea
# 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**: @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.  

---

## 0 answers

