CNN with Two Losses

wwaayyaaww: <p>Hi.I am having a CNN with two CE losses (due to two separate datasets to be trainined simultaneously), each with an individual FC layer prior to the last CE entropy loss estimation.</p>
<p>What should I do if I want to freeze one of these losses, i.e. training only either one without removing any of them?</p>
<ol>
<li>
<p>I have set the requires_grad of one end to be False.</p>
</li>
<li>
<p>Should I set the corresponding layers to the evaluation mode? and what about the gradient calculated for the one to be freezed, simply set it to be zero?</p>
</li>
</ol>
<p>Thank you.</p>

ptrblck: <p>It depends on your current workflow.<br>
I.e. are you using both linear layers, even though are are currently only using a single <code>Dataset</code>?<br>
If so, then setting <code>requires_grad=False</code> for the linear layer, which should not be used, should be enough.<br>
Calling <code>eval()</code> on it shouldn’t be necessary, if you are only using a linear layer, but won’t hurt.</p>

wwaayyaaww: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="77432">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>n though are are currently only using a single <code>Dataset</code> ?<br>
If so, then setting <code>requires_grad=False</code> for the linear layer,</p>
</blockquote>
</aside>
<p>Thank you for your response.</p>
<p>I am having a resnet appended with a linear FC layer for dataset 1, say FC-Dataset-1.</p>
<pre><code class="lang-auto">resnet = ...
in_fea = resnet.logits.in_features
resnet.logits = nn.Linear(in_features = in_fea, out_features = num_sub_1, bias = True)
</code></pre>
<p>In the meantime, there is another linear FC layer for dataset 2, say FC-Dataset-2</p>
<pre><code class="lang-auto">new_logits = nn.Linear(in_features = in_feas, out_features = num_sub_2, bias = True)
</code></pre>
<p>To train the new_logits only (while freezing FC-Dataset-1), I simply set the .requires_grad for FC-Dataset-1 to zero as follows.</p>
<pre><code class="lang-auto">for name, param in resnet.named_parameters():
        if name in ['logits.weight', 'logits.bias']:
            param.requires_grad = False
        else:
            param.requires_grad = True
</code></pre>
<p>And during training, I neglected the output for FC-Dataset-1 as follows.</p>
<pre><code class="lang-auto"># x, y_pred_1 = resnet(x, y)
x, _ = resnet(x, y) 
y_pred_2 = new_logits(x)

# loss_1 = loss_fn_1(y_pred_1, y)
loss_1 = 0
loss_2 = loss_fn_2(y_pred_2, y)
loss = loss_1 + loss_2

#  Involved only 1 optimizer combining parameters for both ends        
optimizer.zero_grad()  
loss.backward()
optimizer.step()
</code></pre>
<p>Is my implementation reasonable? thank you very much.</p>

ptrblck: <p>I assume you’ve written a custom <code>resnet</code> model, as you are expecting two outputs?</p>
<p>Note that your current approach will not switch between the linear layers, but call <code>new_logits</code> on top of <code>resnet.logits</code>. Also, <code>y_pred1</code> isn’t defined in your code.</p>
<p>If you haven’t defined a custom resnet but are using the <code>torchvision</code> implementation, note that this model does not contain a <code>resnet.logits</code> layer and you are simply assigning a new linear layer to this attribute.</p>

wwaayyaaww: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="77432">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Note that your current approach will not switch between the linear layers, but call <code>new_logits</code> on top of <code>resnet.logits</code> . Also, <code>y_pred1</code> isn’t defined in your code.</p>
</blockquote>
</aside>
<p>Yes its a custom resnet model.</p>
<p>The main reason why I do not define <code>y_pred_1</code> is that both dataset 1 and dataset 2 are having different subjects.</p>
<p>Since the inputs are images of dataset 2, defining <code>y_pred_1</code> gives an Cuda error due to subject mismatched.</p>
<pre><code class="lang-auto">Note that your current approach will not switch between the linear layers,
</code></pre>
<p>What would you suggest if I am to switch between the two linear layers?</p>

ptrblck: <aside class="quote no-group" data-username="wwaayyaaww" data-post="5" data-topic="77432">
<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/wwaayyaaww/48/22536_2.png" class="avatar"> wwaayyaaww:</div>
<blockquote>
<p>What would you suggest if I am to switch between the two linear layers?</p>
</blockquote>
</aside>
<p>One approach would be to pass a flag to the <code>forward</code> method and use it as a condition to switch between the linear heads in your model.</p>

wwaayyaaww: <p>Sorry I just revised the snippet earlier. The following is my actual implementation of which I trained only one end while freezing the another.</p>
<pre><code class="lang-auto"># x, y_pred_1 = resnet(x, y)
x, _ = resnet(x, y) 
y_pred_2 = new_logits(x)

# loss_1 = loss_fn_1(y_pred_1, y)
loss_1 = 0
loss_2 = loss_fn_2(y_pred_2, y)
loss = loss_1 + loss_2`
</code></pre>

ptrblck: <p>Thanks for the update.<br>
It looks generally alright, but you are still using this approach (assuming <code>y</code> is a sample from <code>Dataset2</code>):</p>
<pre><code class="lang-python">x -&gt; resnet -&gt; output -&gt; criterion -&gt; loss1
y -&gt; resnet -&gt; output -&gt; new_logits -&gt; criterion -&gt; loss2
</code></pre>
<p>Is that your workflow or would you rather want to switch internally between the last linear layer?</p>

wwaayyaaww: <p>Hmmmm actually I am pre-training the net using  <code>Dataset2</code> , and the entire net will be activated to train both losses with the two datasets simultaneously after this.</p>

ptrblck: <p>Does it mean you would remove <code>new_logits</code> after the pretraining is done?<br>
If that’s the case, then your model seems to be fine.</p>

wwaayyaaww: <p>No, both losses (resnet.logits and new_logits) will be trained together, with two inputs from each dataset 1 and 2 .</p>