Conditionally backpropagating the loss

username: Qinru_Li
cooked: <p>Say in the forward() pass, I generate the network predictions like the following</p>
<pre><code class="lang-auto">def forward(self, input):
    conditions = [1, 2, 3, 4, 5]
    
    out_dict = {}
    for c in conditions:
        output = self.network( (input, c) ) 
        out_dict[index of c] = output

    return out_dict
</code></pre>
<p>Now from the ground truth, we know that we only care about condition 3, and only want to backpropagate the loss in this condition. So I did</p>
<pre><code class="lang-auto">loss = loss_fn(out_dict[3])
loss.backward()
</code></pre>
<p>To verify this, I add a hook to the output of each condition. The code is shown below. If everything works perfectly, we should see “3 [gradient tensor …]”, but it is always “5 [gradient tensor …]”. So it seems like the conditional backpropagation is not working, what is going on here? and how should I do the backprop only with the tensor that I am interested in?</p>
<pre><code class="lang-auto">def forward(self, input):
    conditions = [1, 2, 3, 4, 5]
    
    out_dict = {}
    for c in conditions:
        output = self.network( (input, c) ) 
       
        output.register_hook(lambda grad: print("{index of c}", grad))

        out_dict[index of c] = output

    return out_dict

</code></pre>

username: ptrblck
cooked: <p>The print statement might be buggy, if you are directly printing <code>c</code> in the lambda expression as explained <a href="https://docs.python.org/3/faq/programming.html#why-do-lambdas-defined-in-a-loop-with-different-values-all-return-the-same-result">here</a>, so you might need to assign the value of <code>c</code> as a new argument to the <code>lambda</code> function via:</p>
<pre><code class="lang-python">lambda grad, c=c: print(c, grad)
</code></pre>

username: Qinru_Li
cooked: <p>Ah! That’s it! Thank you for the solution.<br>
I also create a small sample code to play around with this. As you can see, the lambda printouts the correct index only if I am using a local variable.</p>
<p>A weird thing is that when I try to pre-allocate a string <code>message</code> and pass it into the lambda function, the output is still 5.0. I dug into the internet and did not find out when is the string be evaluated, but I hypothesize that it is only evaluated when the <code>__str()__</code> function is called. Therefore, it is keeping a reference to the <code>cond</code>, which has been set to 5.0 in the end (i.e. the same reason as the lambda function)</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        in_features = 5
        self.networks = nn.Sequential(
            nn.Linear(in_features=in_features, out_features=100, bias=True),
            nn.Linear(in_features=100, out_features=2),
        )
        self.initial_weights()

    def forward(self, input):
        num_conditions = 5

        conditions = torch.arange(start=1, end=num_conditions + 1, dtype=torch.float)
        out_features = {}
        for i, cond in enumerate(conditions):
            feature = torch.cat((input, cond.unsqueeze(0)), dim=0)
            feature = self.networks(feature)

            # message = "condition {}".format(cond)
            feature.register_hook(lambda grad, c=cond: print(c))
            out_features[i + 1] = feature
        return out_features

    def initial_weights(self):
        for module in self.networks:
            if isinstance(module, nn.Linear):
                if module.weight is not None:
                    nn.init.xavier_uniform_(module.weight, gain=1.0)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)


torch.manual_seed(42)
network = MyNet()
gt_cond = 3
gt = torch.ones(2)
data = torch.rand(4)
out_features = network(data)

loss = nn.functional.mse_loss(gt, out_features[gt_cond])
loss.backward()  # Expect to see 3.0 in the stdout
</code></pre>