Pytorch `bachward()` updates multiple models

user_1958160 0
Viewed 123

Can anyone tell me why the gradients of the discriminator change as well and if there is a way to avoid it?

for i in range(2):

    X_fake = gen_model(z)

    pred_real = disc_model(X)
    pred_fake = disc_model(X_fake.detach())
    disc_loss = (loss_fn(pred_real, y) + loss_fn(pred_fake, y)) / 2

    disc_optimizer.zero_grad()
    disc_loss.backward()
    disc_optimizer.step()

    pred_fake = disc_model(X_fake)
    gen_loss = loss_fn(pred_fake, y)

    gen_optimizer.zero_grad()
    i == 1 and print_grads(disc_model) # Checkpoint 1
    gen_loss.backward()
    i == 1 and print_grads(disc_model) # Checkpoint 2
    gen_optimizer.step()

This is the rest of the code.

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self._linear = nn.Sequential( nn.Linear(1, 5) )

    def forward(self, X):
        return self._linear(X)

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self._linear = nn.Sequential( nn.Linear(5, 1) )

    def forward(self, X):
        return self._linear(X)

def print_grads(model):
    for params in model.parameters():
        print(params.grad)

# Build the model and data.
gen_model = Generator()
gen_optimizer = torch.optim.Adam(gen_model.parameters(), 1)

disc_model = Discriminator()
disc_optimizer = torch.optim.Adam(disc_model.parameters(), 1)

loss_fn = torch.nn.BCEWithLogitsLoss()

z = torch.rand((1, 1))
X = torch.rand((1, 5))
y = torch.rand((1, 1))
pytorch python
1 Answers

The gradients of the discriminator have been updated because you have backpropagated the loss gen_loss through the discriminator up to the generator itself. In your current formulation of the training loop, this is not a problem since only the generator's parameters will get updated on the next line with gen_optimizer.step(). In other words: yes the gradients of the discriminator will be updated (they change from #1 to #2) but these won't be used to update the parameters of the discriminator. As long as you properly clear the gradient cache with zero_grad before backpropagating, you will be fine.

A common practice is to freeze the discriminator when training the generator, this avoids unnecessary gradient computation:

# freeze discriminator
disc_model.requires_grad_(False)

# train the generator
gen_optimizer.zero_grad()
gen_loss.backward()
gen_optimizer.step()

# unfreeze discriminator for next step
disc_model.requires_grad_(True)
user_63313690
Related