ageron / ageron/handson-ml3

[bug] Training loop error in GANs

Open
#252 0 comments 0 reactions 1 assignee Claimed by @ageron View on GitHub
bug help wanted
Dominant language
Jupyter Notebook
Stars
14.1k
Forks
5.3k
PR merge metrics
No merged PRs in 30d

Description

### Enter the chapter number

Chapter-17: Autoencoders, GANs and Diffusion Models

### Enter the page number

_No response_

### What is the cell's number in the notebook

Cell 54

### Enter the environment you are using to run the notebook

Kaggle

### Describe your issue

Running the training loop gives
```python
UserWarning: The model does not have any trainable weights. warnings.warn("The model does not have any trainable weights.")
```
error, the error is due to the updated Keras API and the discriminator and generator needs to be updated using tf.GradientTape() to override the train_step

### Enter what you expected to happen

_No response_

### If you found a workaround, describe it here

Creating a custom GAN class would be better here,
```python
class GAN(keras.Model):
def __init__(self, discriminator, generator, codings_size):
super().__init__()
self.discriminator = discriminator
self.generator = generator
self.codings_size = codings_size

self.d_loss_tracker = tf.keras.metrics.Mean(name="d_loss")
self.g_loss_tracker = tf.keras.metrics.Mean(name="g_loss")

@property
def metrics(self):
return [self.d_loss_tracker, self.g_loss_tracker]

def compile(self, d_optimizer, g_optimizer, loss_fn):
super().compile()
self.d_optimizer = d_optimizer
self.g_optimizer = g_optimizer
self.loss_fn = loss_fn

def train_step(self, real_images):
batch_size = tf.shape(real_images)[0]

# Train the discriminator

random_latent_vectors = tf.random.normal(shape=(batch_size, self.codings_size))
generated_images = self.generator(random_latent_vectors, training=True)
combined_images = tf.concat([generated_images, real_images], axis=0)

labels = tf.concat([tf.zeros((batch_size, 1)), tf.ones((batch_size, 1))], axis=0)

with tf.GradientTape() as tape:
predictions = self.discriminator(combined_images, training=True)
d_loss = self.loss_fn(labels, predictions)

d_gradients = tape.gradient(d_loss, self.discriminator.trainable_weights)
self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_weights))

# Train the generator

random_latent_vectors = tf.random.normal(shape=(batch_size, self.codings_size))
misleading_labels = tf.ones((batch_size, 1))

with tf.GradientTape() as tape:
generated_images = self.generator(random_latent_vectors, training=True)
predictions = self.discriminator(generated_images, training=True)
g_loss = self.loss_fn(misleading_labels, predictions)

g_gradients = tape.gradient(g_loss, self.generator.trainable_weights)
self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_weights))

self.d_loss_tracker.update_state(d_loss)
self.g_loss_tracker.update_state(g_loss)

return {"d_loss": self.d_loss_tracker.result(), "g_loss": self.g_loss_tracker.result()}
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.