How to disable discriminator training inside a Tensorflow GAN?

Question

I have a generator like this:

def make_generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(7*7*256, use_bias=False, input_shape=(100,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 256)))
    assert model.output_shape == (None, 7, 7, 256)  # Note: None is the batch size

    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same', use_bias=False))
    assert model.output_shape == (None, 7, 7, 128)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same', use_bias=False))
    assert model.output_shape == (None, 14, 14, 64)
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False, activation='tanh'))
    assert model.output_shape == (None, 28, 28, 1)

    return model

And a discriminator:

def make_discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',
                                     input_shape=[28, 28, 1]))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.4))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.4))

    model.add(layers.Flatten())
    model.add(layers.Dense(1, activation='sigmoid'))
    opt = tf.keras.optimizers.Adam(lr=0.0002, beta_1=0.5)
    model.compile(loss='binary_crossentropy', optimizer=opt, metrics=['accuracy'])
    return model

I group them inside a GAN:

def make_gan(g_model, d_model):
    d_model.trainable = False #<-- Not working?
    model = tf.keras.Sequential()
    model.add(g_model)
    model.add(d_model)
    print("Trainable GAN params before:" + str(model.trainable_variables))
    d_model.trainable = False # <-- Also not working?
    print("Trainable GAN params after:" + str(model.trainable_variables)) #<-- The same as before
    opt = tf.keras.optimizers.Adam(lr=0.0002, beta_1=0.5)
    model.compile(loss='binary_crossentropy', optimizer=opt)
    return model

For the GAN training to work, I need to train the Generator and Discriminator separately. I train the generator, using the GAN model, but I want to lock the discriminator part in place during training. This doesn't seem to work, as my discriminator always ends up classifying everything as real within 1 or 2 epochs. This leads to my GAN not making any progress after around 50 epochs.

Why does d_model.trainable=False not work in this context?

Here is my Train loop:

def train(g_model, d_model, gan_model, dataset, latent_dim, n_epochs=100, n_batch=256):
    batch_per_epoch = int(dataset.shape[0] / n_batch)
    half_batch = int(n_batch / 2)
    for i in range(n_epochs):
        #train discriminator:
        for j in range(batch_per_epoch):
            X_real, y_real = generate_real_samples(dataset, half_batch)
            X_fake, y_fake = generate_fake_samples(g_model, latent_dim, half_batch)
            X, y = vstack((X_real, X_fake)), vstack((y_real, y_fake))
            d_fake_loss, _ = d_model.train_on_batch(X_real, y_real)
            d_real_loss, _ = d_model.train_on_batch(X_fake, y_fake)
            if j % 10 == 0:
                print('>Epoch %d - Training Discriminator epoch %d: %d/%d, d_real_loss=%.3f' % (i+1, k+1, j+1, batch_per_epoch, d_real_loss))
                    print('>Epoch %d - Training Discriminator epoch %d: %d/%d, d_fake_loss=%.3f' % (i+1, k+1, j+1, batch_per_epoch, d_fake_loss))
                    #print('>Epoch %d - Training Discriminator epoch %d: %d/%d, d_loss=%.3f' % (i+1, k+1, j+1, batch_per_epoch, d_loss))
        #train generator
        for j in range(batch_per_epoch):        
            X_gan = generate_latent_points(latent_dim, n_batch)
            y_gan = ones((n_batch, 1))
            g_loss = gan_model.train_on_batch(X_gan, y_gan)
            if j % 10 == 0:
                print('>Epoch %d - Training Generator: %d/%d, g_loss=%.3f' % (i+1, j+1, batch_per_epoch, g_loss))
        #call the generator to generate an image end the end of each epoch
        #x_generated, _ = generate_fake_samples(g_model, latent_dim, 100)
        #save_plot(x_generated, i)
        summarize_performance(i, g_model, d_model, dataset, latent_dim)

user_17905764 · Answer

in tensorflow, there is no need to explicitly lock the weights. You do the following as shown in the code:

Within the GradientTape you propagate your data through the generator, pass a real image and a generated image to the discriminator and finally calculate the loss
get the gradients with respect to the loss. Here you pass only the trainable weights of your respective model.

Finally, apply the gradients you calculated to the respective model.

with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
  generated_images = self.generator(noise, training=True)

  real_output = self.discriminator(batch, training=True)
  fake_output = self.discriminator(generated_images, training=True)

  gen_loss = self.generator_loss(fake_output)
  disc_loss = self.discriminator_loss(real_output, fake_output)

gradients_of_generator = gen_tape.gradient(gen_loss, self.generator.trainable_variables)
gradients_of_discriminator = disc_tape.gradient(disc_loss, self.discriminator.trainable_variables)

self.generator_optimizer.apply_gradients(zip(gradients_of_generator, self.generator.trainable_variables))
self.discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, self.discriminator.trainable_variables))
return [gen_loss, disc_loss]

In this example, two gradient tapes are used within the scope of with tf.GradientTape. The reason is, that memory of a tape is released as soon as the values are retrived with tape.gradient(). You could also use a single tape by using with tf.GradientTape(persistent=True) as gradient_tape:

How to disable discriminator training inside a Tensorflow GAN?

1 Answers