How to prevent one output in a multi-output network?

Jingnan_Jia: <p>My code for multi-task training (segmentation task and reconstruction task):</p>
<pre><code class="lang-auto">class MyNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.layer1 = nn.Conv2d(3, 32, 3)
    self.layer2 = nn.Conv2d(32, 64, 3)
    self.layer3 = nn.Conv2d(64, 2, 3)
    self.layer4 = nn.Conv2d(64, 1, 3)

  def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)
    out1 = self.layer3(out)
    out2 = self.layer4(out)

    return out1, out2


</code></pre>
<p>During training, I get 2 outputs and compute their losses respectively. So that I can train the net alternatively.</p>
<pre><code class="lang-auto">x, y = dataloader()  # generate a batch of data
net = MyNet()
opt = torch.optim.Adam(self.net.parameters(), lr=1e-4)

pred1, pred2 = net(x)
loss1 = seg_loss(pred1, y)
loss2 = rec_loss(pred2, x)
loss1.backward()
opt.step()
loss2.backward()
opt.step()
```
Correct, right?

One day, I want to just train the reconstruction branch. The I would use the following code:
</code></pre>
<p>x, y = dataloader()  # generate a batch of data<br>
net = MyNet()<br>
opt = torch.optim.Adam(self.net.parameters(), lr=1e-4)</p>
<p>pred1, pred2 = net(x)<br>
loss2 = rec_loss(pred2, x)<br>
loss2.backward()<br>
opt.step()</p>
<pre><code class="lang-auto">
It works well. Pred1 does not participate backward as I expected. But I found that the forward of pred1 also accupy some GPU memory. In order to save GPU memory, how to prevent the forward of pred1???</code></pre>

t_naumenko: <p>If you want to keep this output during reconstruction-only training, and prevent autograd from storing its intermediate tensors, you can wrap <code>self.layer3(out)</code> execution with <code>torch.set_grad_enabled(False)</code> context manager. In order to swap between your training regimes, you can use the following code example:</p>
<pre><code class="lang-auto">class MyNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.seg_train = True
    self.layer1 = nn.Conv2d(3, 32, 3)
    self.layer2 = nn.Conv2d(32, 64, 3)
    self.layer3 = nn.Conv2d(64, 2, 3)
    self.layer4 = nn.Conv2d(64, 1, 3)

  def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)

    # in grad is not enabled - we should not enable it
    grad_enabled = torch.is_grad_enabled()
    with torch.set_grad_enabled(self.seg_train and grad_enabled):
        out1 = self.layer3(out)

    out2 = self.layer4(out)

    return out1, out2
</code></pre>
<p>Then you will be able to swap training regimes by chaning <code>model.seg_train</code> variable.</p>

Jingnan_Jia: <p>thanks, another question is:</p>
<p>If I finish the training of “Seg+Recon” Net, and save the whole net to disk. In the future if I want to load the trainned net and continue to fine-tune the “Recon” Net. Can your code work in this case??</p>
<p>Or If I finish the training of “Recon” Net, and save the “Recon” Net to disk. In the future if I want to load the trainned “Recon” and continue to fine-tune the “Seg+Recon” Net. Can your code work in this case??</p>
<p>I am just curious if the saved model will include the “self.layer3” even if I did not use it in “forward” function.</p>