How to skip a part of the model for back propagation?

username: werthers
cooked: <p>Hi there,<br>
I’m new to pytorch and I’d like to implement a GAN model, I see there’s a tutorial <a href="https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html" rel="noopener nofollow ugc">DCGAN Tutorial — PyTorch Tutorials 1.12.1+cu102 documentation</a>, it uses two optimizer to backward the parameters of generator and discriminator separately, so that the gradient of the generator loss function is not back propagated into the discriminator parameters, and the gradient of the discriminator loss function is not back propagated into the generator parameters.</p>
<p>But what if I want to build a more complicated model which need to share some lower layers between generator and discriminator, I can’t put the parameters of the generator and the discriminator into different optimizers because they share many parameters. How can I train such a model correctly and efficiently?</p>
<p>Here’s the sample network structure:</p>
<blockquote>
<p>class Generator(nn.Module):<br>
self.encoder: shared network between generator and discriminator<br>
self.top: linear layer</p>
<pre><code>def forward(self, input):
    return self.top(self.encoder(input))
</code></pre>
<p>class Discriminator(nn.Module):<br>
self.encoder: shared network between generator and discriminator<br>
self.top: linear layer</p>
<pre><code>def forward(self, input):
    return self.top(self.encoder(input))
</code></pre>
<p>netG = Generator()<br>
netD = Discriminator()</p>
<p>training loop:<br>
############################<br>
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))<br>
###########################<br>
## Train with all-real batchoutput = netD(real_cpu).view(-1)<br>
# Calculate loss on all-real batch<br>
errD_real = criterion(netD(real_img), real_label)</p>
<pre><code>## Train with all-fake batch
# Calculate D's loss on the all-fake batch
errD_fake = criterion(netD(netG(noise).detach()), fake_label)

# Calculate gradients for D in backward pass
errD = errD_real + errD_fake
errD.backward()

############################
# (2) Update G network: maximize log(D(G(z)))
###########################
netG.zero_grad()
label.fill_(real_label)  # fake labels are real for generator cost
# Since we just updated D, perform another forward pass of all-fake batch through D
output = netD(fake).view(-1)
# Calculate G's loss based on this output
errG = criterion(netD(netG(noise)), fake_label) #!!! We don't need the gradient of netD with respect to errG, but we still need the gradient of netG with respect to errG and the gradient of netD with respect to errD above, so how to skip back propagation of netD for this loss(under the setting that netG and netD share some lower layers)?
# Calculate gradients for G
errG.backward()
</code></pre>
</blockquote>

username: werthers
cooked: <p>When I compute D_fake loss, the whole forward function is Discriminator_top(Discriminator_encoder(Generator_top(Generator_encoder(input)))), and I only need the gradient of Discriminator_top and Discriminator_encoder.</p>
<p>When I compute G loss, the whole forward function is Discriminator_top(Discriminator_encoder(Generator_top(Generator_encoder(input)))), and I only need the gradient of Generator_top and Generator_encoder.</p>
<p>Discriminator_encoder and Generator_encoder is the same network, so in every training step this shared encoder should receive backward gradient for D_fake loss (backward to Discriminator_encoder) and G loss (backward to Generator_encoder)</p>

username: KFrank
cooked: <p>Hi Werther!</p>
<aside class="quote no-group" data-username="werthers" data-post="1" data-topic="164768" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/w/7993a0/48.png" class="avatar"> werthers:</div>
<blockquote>
<p>how to skip back propagation of netD for this loss(under the setting that netG and netD share some lower layers)?</p>
</blockquote>
</aside>
<p>The best way I can think of doing this (if I correctly understand your use<br>
case) is to break your backpropagation up into two pieces using two calls<br>
to <code>torch.autograd.grad()</code> (rather than a single call to <code>.backward()</code>).</p>
<p>Here are two calls to <code>autograd.grad()</code> applied to a simplified version of<br>
what I think you are trying to do:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; input = torch.tensor ([2.0, 10.0])
&gt;&gt;&gt;
&gt;&gt;&gt; p1 = torch.tensor ([1.0], requires_grad = True)   # separate-parameters version
&gt;&gt;&gt; p2 = torch.tensor ([1.0], requires_grad = True)   # separate-parameters version
&gt;&gt;&gt;
&gt;&gt;&gt; o  = p1 * input + p2 * input**2   # run "model" -- vector output
&gt;&gt;&gt; loss = o.sum()                    # scalar loss
&gt;&gt;&gt; loss
tensor(116., grad_fn=&lt;SumBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; p1.grad
tensor([12.])
&gt;&gt;&gt; p2.grad
tensor([104.])
&gt;&gt;&gt;
&gt;&gt;&gt; p = torch.tensor ([1.0], requires_grad = True)   # shared-parameter version
&gt;&gt;&gt;
&gt;&gt;&gt; o = p * input + p * input**2   # run shared-parameter "model"
&gt;&gt;&gt; loss = o.sum()
&gt;&gt;&gt; loss   # loss is the same
tensor(116., grad_fn=&lt;SumBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; p.grad   # grad gets accumulated for both halves of "model"
tensor([116.])
&gt;&gt;&gt;
&gt;&gt;&gt; p.grad = None
&gt;&gt;&gt;
&gt;&gt;&gt; o1 = p * input           # run shared-parameter "model" in two pieces -- first half
&gt;&gt;&gt; o2 = o1 + p * input**2   # second half
&gt;&gt;&gt; loss = o2.sum()
&gt;&gt;&gt; loss   # loss is the same
tensor(116., grad_fn=&lt;SumBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; grad2 = torch.autograd.grad (loss, o1)       # gradient of second half of "model" with respect to o1
&gt;&gt;&gt; grad2    # grad2, as expected, is not a scalar
(tensor([1., 1.]),)
&gt;&gt;&gt; p.grad   # doesn't populate p.grad
&gt;&gt;&gt; grad1 = torch.autograd.grad (o1, p, grad2)   # gradient of the first half of "model" with respect to p
&gt;&gt;&gt; grad1    # same as p1.grad
(tensor([12.]),)
</code></pre>
<p>As an aside, I’m not convinced that it makes sense to share layers between<br>
your generator and discriminator.  The two networks are doing rather different<br>
things, and although it may well make sense for them to both have a layer<br>
of the same architecture, it’s not clear to me that such layers should have<br>
the same values for their parameters.</p>
<p>One could argue that the “features” your <code>encoder</code> layer learns to produce<br>
make sense for both the generation and discrimination process, and that<br>
it is helpful to  train the same set of parameters for both, favoring sharing<br>
the <code>encoder</code>.</p>
<p>But, typically, you input some sort of random noise into your generator, from<br>
which I assume the generator’s <code>encoder</code> produces features, whereas you<br>
input a structured image (or whatever the samples are) – either real or fake,<br>
but still not random noise – into the discriminator’s <code>encoder</code>.  Even if you<br>
want the <code>encoder</code> to produce similar features in both cases, my intuition tells<br>
me that the parameters you would use for producing features from random<br>
noise are likely to be quite different than those your would use for producing<br>
(similar) features from a structured image.</p>
<p>I’m sufficiently skeptical of the shared-encoder idea that I would want to see<br>
a comparison of the shared-encoder approach with the separate-encoder<br>
approach that shows that the former actually works better.</p>
<p>Best.</p>
<p>K. Frank</p>

username: werthers
cooked: <p>Thank you, Frank, that’s exactly the case. Yes, I am doing a study to verify the effect of shared generator, discriminator representations using an encoder with huge capacity. So that’s why I would like to use a correct and efficient way to break backpropagation up into two pieces, I used to copy a frozen discriminator when I backward generator loss, but I think it may not be efficient enough.</p>
<p>I have not used the autograd.grad() method directly before, your solution inspired me a lot. And as it’s a big encoder and huge amount of data, I usually train with the AMP module, so I wonder if this solution suitable for use with AMP module?</p>
<blockquote>
<p>o1 = p * input           # run shared-parameter “model” in two pieces – first half<br>
o2 = o1 + p * input**2   # second half<br>
loss = o2.sum()</p>
<p>grad2 = torch.autograd.grad(loss, o1)       # gradient of second half of “model” with respect to o1</p>
<p>grad1 = torch.autograd.grad(o1, p, grad2)   # gradient of the first half of “model” with respect to p</p>
</blockquote>
<p>Is it safe to convert the above codes to the following ones?</p>
<blockquote>
<p>scaler = torch.cuda.amp.GradScaler()<br>
with torch.cuda.amp.autocast():<br>
o1 = p * input           # run shared-parameter “model” in two pieces – first half<br>
o2 = o1 + p * input**2   # second half<br>
loss = o2.sum()</p>
<p>grad2 = torch.autograd.grad(scaler.scale(loss), o1)       # gradient of second half of “model” with respect to o1</p>
<p>grad1 = torch.autograd.grad(o1, p, grad2)   # gradient of the first half of “model” with respect to p</p>
</blockquote>
<p>Thank you for your help!</p>
<p>Gratefully,</p>
<p>Werther</p>