How to train specific layers with different optimizers

GSauce: <p>I am using an autoencoder,</p>
<pre data-code-wrap="class"><code class="lang-nohighlight">   def __init__(self):
       super(autoencoder, self).__init__()
       #self.channel = channel
       self.encoder = nn.Sequential(
           nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, stride = 1), # 26
           nn.ReLU(),
           nn.MaxPool2d(2, stride = 2), # 13
           nn.Conv2d(in_channels = 32, out_channels = 16, kernel_size = 3, stride = 1), # 11
           nn.ReLU(),
           nn.MaxPool2d(2, stride = 2), # 5
           nn.Conv2d(in_channels = 16, out_channels = 8, kernel_size = 3, stride = 1), # 3
           nn.ReLU(),
       ) # 8x3x3
   
       self.decoder = nn.Sequential(
           nn.ConvTranspose2d(in_channels = 8, out_channels = 16, kernel_size = 3, stride = 1), #5
           nn.ReLU(),
           nn.ConvTranspose2d(in_channels = 16, out_channels = 32, kernel_size = 3, stride = 1), #7
           nn.ReLU(),
           nn.ConvTranspose2d(in_channels = 32, out_channels = 32, kernel_size = 3, stride = 2), #15
           nn.ReLU(),
           nn.ConvTranspose2d(in_channels= 32, out_channels = 3, kernel_size = 3, stride = 2,padding=2, output_padding=1), #28
           nn.Tanh()
       )


       self.classifier = nn.Sequential(
           # nn.AdaptiveAvgPool2d((1,1)),
           nn.Flatten(),
           nn.Linear(8*3*3,8),
           nn.ReLU(),
           nn.Linear(8,10)
       )

   def forward(self, x):
       x= self.encoder(x)
       out = self.classifier(x)
       # feature = self.flatten(x)
       x = self.decoder(x)
       
       return  x,out
</code></pre>
<p>now i want to update the weights of the autoencoder without the classifier using adam optimizer and calculating the <code>L2 loss</code> between the reconstruction images and the original image.<br>
and want to train only the classifier using adam only, but using different loss i.e, <code>crossentropyloss</code>. now how can I train encoder, decoder, and the classifier separately in one go?</p>
<p>I want to do something like this:</p>
<pre><code class="lang-auto">for epoch in range(10):
    print(f"Epoch {epoch+1}")
    
    for images in tqdm(trainloader):
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        
        images,Ys = images
        images = images.cuda()
        Ys = Ys.cuda()
        x, out = model(images)
        
        loss = loss_fn1(x, images) # l2 norm
        loss1 = loss_fn2(out, Ys) # cross entropy
        total_loss+=loss
        total_loss1+=loss1
        
        # print('Loss b/w images: ',loss.item())
        # print('Loss b/w labels: ',loss1.item())
        # break
        loss.backward() # l2 norm
        loss1.backward() # cross en

        optimizer1.step()
        optimizer2.step()
    
</code></pre>
<p>But it is showing error, which I fixed by setting <code>retrain_graph=True</code>, but that was not feasible since it is taking all my GPU memory. Is there any other alternative?</p>

KFrank: <p>Hi GSauce!</p>
<aside class="quote no-group" data-username="GSauce" data-post="1" data-topic="162748" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gsauce/48/51727_2.png" class="avatar"> GSauce:</div>
<blockquote>
<pre><code class="lang-auto">    for images in tqdm(trainloader):
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        
        images,Ys = images
        images = images.cuda()
        Ys = Ys.cuda()
        x, out = model(images)
        
        loss = loss_fn1(x, images) # l2 norm
        loss1 = loss_fn2(out, Ys) # cross entropy
        total_loss+=loss
        total_loss1+=loss1
        
        # print('Loss b/w images: ',loss.item())
        # print('Loss b/w labels: ',loss1.item())
        # break
        loss.backward() # l2 norm
        loss1.backward() # cross en

        optimizer1.step()
        optimizer2.step()    
</code></pre>
</blockquote>
</aside>
<p>As written, I don’t believe that this does what you want.  <code>loss.backward()</code><br>
<em>accumulates</em> gradients, so your two <code>.backward()</code> calls will accumulate<br>
gradients from both of your loss functions into the parameters of both of<br>
your optimizers.</p>
<p>If I understand your use case, you want something like:</p>
<pre data-code-wrap="python"><code class="lang-python">        loss = loss_fn1(x, images) # l2 norm
        loss1 = loss_fn2(out, Ys) # cross entropy

        optimizer1.zero_grad()
        loss.backward (retain_graph = True)  # l2 norm
        optimizer1.step()   # only takes a step with loss gradients

        optimizer2.zero_grad()   # zeroes out gradients from loss
        loss1.backward() # cross en
        optimizer2.step()   # only takes a step with loss1 gradients
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>which I fixed by setting <code>retrain_graph=True</code></p>
</blockquote>
</aside>
<p>Yes, as a general rule if you call <code>backward()</code> more than once for the same<br>
forward pass, you will need to use <code>retain_graph = True</code> for all but the<br>
last call.</p>
<aside class="quote no-group">
<blockquote>
<p>but that was not feasible since it is taking all my GPU memory.</p>
</blockquote>
</aside>
<p>And yes, retaining the graph does keep the graph in memory.</p>
<p>This <a href="https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" rel="noopener nofollow ugc">performance-tuning guide</a> does mention some things you can do<br>
to reduce memory usage.</p>
<p>Also, the following thread discusses a use case where some parts of<br>
a shared computation graph could be released “early” while keep other<br>
parts that were necessary for a subsequent <code>.backward()</code> call.  It may<br>
or may not be relevant to your use case, but it might be worth a look.</p>
<aside class="quote quote-modified" data-post="1" data-topic="154885">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/leanderk/48/8099_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/eager-autograd-on-leaves/154885">Eager Autograd on Leaves</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    I am working on a huge model where I am very much memory constrained. I continuously generate terms for regularisation, my loss looks something like: 
loss = accuracy + sum_ penalty alpha*penalty 
my computational graph looks like this: 
computation -&gt; computation -&gt; computation -&gt; computation -&gt; computation -&gt; (...)
                    |                          |                      
                    -&gt; penalty                 -&gt; penalty                                                     …
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>

GSauce: <p>Thank you for replying, okay understood, I tried one more thing,</p>
<pre><code class="lang-auto">l = loss + loss1
l.backward()
optim1.step()
optim2.step()
</code></pre>
<pre><code class="lang-python">model = autoencoder().cuda()

parameters = list()
parameters.extend(model.decoder.parameters())
parameters.extend(model.encoder.parameters())

loss_fn1 = nn.MSELoss()
loss_fn2 = nn.CrossEntropyLoss()
optimizer1 = optim.Adam(parameters, lr=lr)
optimizer2 = optim.Adam(model.parameters(), lr=3e-4)

    for images in tqdm(trainloader):
        optimizer1.zero_grad()
        optimizer2.zero_grad()
        
        images,Ys = images
        images = images.cuda()
        Ys = Ys.cuda()
        x, out = model(images)
        
        loss = loss_fn1(x, images) # l2 norm
        loss1 = loss_fn2(out, Ys) # cross entropy
        
        l = loss+loss1
        
        l.backward() 
        total_loss+=l
        total_loss1+=loss1
        optimizer1.step()
        optimizer2.step()
</code></pre>
<p>this is roughly what I did. but I had one doubt, that will <code>.backward()</code> generates a single graph even though there are 2 loss functions?</p>

srishti-git1110: <p>Yes, there will be one backward (and one forward) graph for the tensor <code>l</code>.</p>

srishti-git1110: <aside class="quote no-group" data-username="GSauce" data-post="1" data-topic="162748">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gsauce/48/51727_2.png" class="avatar"> GSauce:</div>
<blockquote>
<p>now i want to update the weights of the autoencoder without the classifier using adam optimizer and calculating the <code>L2 loss</code> between the reconstruction images and the original image.<br>
and want to train only the classifier using adam only, but using different loss i.e, <code>crossentropyloss</code>.</p>
</blockquote>
</aside>
<p>Also note that what you are willing to do here isn’t what’s exactly being done with this code of yours:</p>
<aside class="quote no-group" data-username="GSauce" data-post="3" data-topic="162748">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gsauce/48/51727_2.png" class="avatar"> GSauce:</div>
<blockquote>
<pre><code class="lang-auto">l = loss+loss1
        
        l.backward() 
        total_loss+=l
        total_loss1+=loss1
        optimizer1.step()
        optimizer2.step()
</code></pre>
</blockquote>
</aside>
<p>Here, <code>l.backward()</code> shall calculate gradients for the model parameters based on both <code>loss</code> and <code>loss1</code> as opposed to what you stated in the first post.</p>

GSauce: <p>Thank you for your response.<br>
Exactly, that is what I was thinking, I tried doing that using <code>retain_graph=True</code> at expense of my GPU memory, and I wasn’t able to train due to lack of memory. which is why I tried this one out, after googling.</p>
<p>what I want to achieve is, the entire <strong>encoder and decoder</strong> to update their weights based on the MSELoss calculated by comparing each and every pixel of the input image and the reconstruction image.<br>
But I want to simultaneously train the classifier based on CrossEntropyLoss, and I am using a different optimizer to train this.</p>
<p>NOTE:: the 2nd optimizer should not change the weights of the <strong>encoder or decoder</strong> it should only update the classifier.</p>
<p>So this is what I want to achieve but still couldn’t find an efficient way.</p>

srishti-git1110: <aside class="quote no-group" data-username="GSauce" data-post="6" data-topic="162748">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/gsauce/48/51727_2.png" class="avatar"> GSauce:</div>
<blockquote>
<p>the 2nd optimizer should not change the weights of the <strong>encoder or decoder</strong> it should only update the classifier.</p>
</blockquote>
</aside>
<p>In this case, change optimizer2 as:</p>
<pre><code class="lang-auto">optimizer1 = optim.Adam(parameters, lr=lr)  # encoder and decoder parameters
optimizer2 = optim.Adam(model.classifier.parameters(), lr=3e-4)  # parameters from only the classifier layer
</code></pre>
<p>Still, keep in mind that <code>l.backward()</code> will calculate the gradients based on both <code>loss</code> and <code>loss1</code>.<br>
The only straight forward solution I could think of is what <a class="mention" href="/u/kfrank">@KFrank</a> has coded in their post. Did you check out the tuning guide they linked for your memory issues?</p>