Loss function backpropagation

ShanakaRG: <p>Hi all,<br>
I have a two stream network like this</p>
<p>import torch<br>
import torch.nn as nn<br>
import torch.optim as optim<br>
from torchviz import make_dot</p>
<h1>
<a name="define-the-two-stream-network-with-a-shared-layer-1" class="anchor" href="#define-the-two-stream-network-with-a-shared-layer-1"></a>Define the two-stream network with a shared layer</h1>
<p>class TwoStreamNet(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(TwoStreamNet, self).<strong>init</strong>()<br>
self.shared_layer = nn.Sequential(<br>
nn.Linear(10, 20),<br>
nn.ReLU()<br>
)<br>
self.stream1 = nn.Sequential(<br>
nn.Linear(20, 10),<br>
nn.ReLU(),<br>
nn.Linear(10, 5),<br>
nn.ReLU()<br>
)<br>
self.stream2 = nn.Sequential(<br>
nn.Linear(20, 10),<br>
nn.ReLU(),<br>
nn.Linear(10, 5),<br>
nn.ReLU()<br>
)</p>
<pre><code>def forward(self, x):
    out_shared = self.shared_layer(x)
    out1 = self.stream1(out_shared)
    out2 = self.stream2(out_shared)
    return out1, out2
</code></pre>
<h1>
<a name="create-a-sample-input-tensor-2" class="anchor" href="#create-a-sample-input-tensor-2"></a>Create a sample input tensor</h1>
<p>x = torch.randn(1, 10)</p>
<h1>
<a name="initialize-the-network-3" class="anchor" href="#initialize-the-network-3"></a>Initialize the network</h1>
<p>net = TwoStreamNet()</p>
<h1>
<a name="compute-the-output-of-the-network-4" class="anchor" href="#compute-the-output-of-the-network-4"></a>Compute the output of the network</h1>
<p>out1, out2 = net(x)</p>
<h1>
<a name="define-the-loss-function-5" class="anchor" href="#define-the-loss-function-5"></a>Define the loss function</h1>
<p>loss_fn1 = nn.CrossEntropyLoss()<br>
loss_fn2 = nn.CrossEntropyLoss()<br>
loss1 = loss_fn1(out1, torch.ones_like(out1))<br>
loss2 = loss_fn2(out2, torch.zeros_like(out2))<br>
loss =  loss1*loss2</p>
<h1>
<a name="backpropagate-the-loss-6" class="anchor" href="#backpropagate-the-loss-6"></a>Backpropagate the loss</h1>
<p>loss.backward(retain_graph=True)</p>
<h1>
<a name="visualize-the-gradients-for-the-weights-in-stream-2-7" class="anchor" href="#visualize-the-gradients-for-the-weights-in-stream-2-7"></a>Visualize the gradients for the weights in stream 2</h1>
<h1>
<a name="printgradients-for-stream-2-8" class="anchor" href="#printgradients-for-stream-2-8"></a>print(‘Gradients for stream 2:’)</h1>
<h1>
<a name="printnetstream20weightgrad-9" class="anchor" href="#printnetstream20weightgrad-9"></a>print(net.stream2[0].weight.grad)</h1>
<p>make_dot((loss1), params=dict(net.named_parameters()))</p>
<p>i need to make sure my loss1 impacts stream 2 in the back propagation. how can i visualize the backward parth of the loss?  i used make_dot to get the graph. but it does not solve my problem because when i visualize the graph belongs to loss1, it does not have any impact with the weights of stream2.<br>
please help me to understand the pytorch backwards function.</p>

ptrblck: <aside class="quote no-group" data-username="ShanakaRG" data-post="1" data-topic="174283">
<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/shanakarg/48/58080_2.png" class="avatar"> ShanakaRG:</div>
<blockquote>
<p>i need to make sure my loss1 impacts stream 2 in the back propagation.</p>
</blockquote>
</aside>
<p><code>loss1</code> won’t directly contribute to the gradient calculation on <code>stream2</code> since no parameters in <code>stream2</code> were used to calculate <code>loss1</code>, so there is no direct dependency.<br>
However, <code>stream2</code> consumes <code>out_shared</code> created by <code>shared_layer</code>, which was used to compute <code>loss1</code>, so its gradients will be affected by <code>loss1</code>.</p>

ShanakaRG: <p>Thank you for the reply <a class="mention" href="/u/ptrblck">@ptrblck</a> !! But what if i  use the total loss = loss1* loss2 for the backpropagation. Then the loss1 should have a contribution to update the weights in the stream2. I wrote the partial derivatives. Please guide me if i did any mistake !! as well as let me know that pytorch backwards can handle this case as i wrote in the picture</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/8/a8771ef62d2f238064b4b6bcc56c7e98a5058799.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a8771ef62d2f238064b4b6bcc56c7e98a5058799" title="pic"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/8/a8771ef62d2f238064b4b6bcc56c7e98a5058799_2_375x500.jpeg" alt="pic" data-base62-sha1="o2jqGLA8dAF2IsXuK96oq4eDyXv" width="375" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/8/a8771ef62d2f238064b4b6bcc56c7e98a5058799_2_375x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/8/a8771ef62d2f238064b4b6bcc56c7e98a5058799_2_562x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/8/a8771ef62d2f238064b4b6bcc56c7e98a5058799_2_750x1000.jpeg 2x" data-dominant-color="D3B4BA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">pic</span><span class="informations">1536×2048 344 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>Yes, you are right. The multiplication would create the interaction as seen in this small example (I thought you were adding the losses together):</p>
<pre><code class="lang-python"># setup
lin1 = nn.Linear(1, 1, bias=False)
lin2 = nn.Linear(1, 1, bias=False)
lin1.zero_grad()
lin2.zero_grad()
x1 = torch.randn(1, 1)
x2 = torch.randn(1, 1)

# independent forward passes
out1 = lin1(x1)
out2 = lin2(x2)
out1.backward()
out2.backward()

print(lin1.weight.grad, x1)
# tensor([[-0.3172]]) tensor([[-0.3172]])
print(lin2.weight.grad, x2)
# tensor([[-0.2574]]) tensor([[-0.2574]])

g1 = lin1.weight.grad.clone()
g2 = lin2.weight.grad.clone()

lin1.zero_grad()
lin2.zero_grad()

# addition
out1 = lin1(x1)
out2 = lin2(x2)

loss = out1 + out2
loss.backward()
print(lin1.weight.grad, g1, x1)
# tensor([[-0.3172]]) tensor([[-0.3172]]) tensor([[-0.3172]])
print(lin2.weight.grad, g2, x2)
# tensor([[-0.2574]]) tensor([[-0.2574]]) tensor([[-0.2574]])

lin1.zero_grad()
lin2.zero_grad()

# multiplication
out1 = lin1(x1)
out2 = lin2(x2)

loss = out1 * out2
loss.backward()
print(lin1.weight.grad, g1, x1*out2)
# tensor([[0.0582]]) tensor([[-0.3172]]) tensor([[0.0582]], grad_fn=&lt;MulBackward0&gt;)
print(lin2.weight.grad, g2, x2*out1)
# tensor([[0.0215]]) tensor([[-0.2574]]) tensor([[0.0215]], grad_fn=&lt;MulBackward0&gt;)
</code></pre>

ShanakaRG: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> .<br>
yes i’m multiplying losses .<br>
so if i use the torchviz to visualize the gradient path as in the following code, does it give the result of backpropagation with respect to loss1 ?<br>
I’m confused because it does not involve any parameter related to stream 2 as shown in the graph.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim
import torch.optim as optim
from torchviz import make_dot
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()
class SharedLayerNet(nn.Module):
    def __init__(self):
        super(SharedLayerNet, self).__init__()
        self.shared_fc = nn.Linear(10, 5) # Shared layer
        self.stream1_fc = nn.Linear(5, 3) # First stream
        self.stream2_fc = nn.Linear(5, 2) # Second stream
        
    def forward(self, x):
        shared_out = self.shared_fc(x) # Shared layer output
        out1 = self.stream1_fc(shared_out) # First stream output
        out2 = self.stream2_fc(shared_out) # Second stream output
        return out1, out2
    
model = SharedLayerNet()
criterion1 = nn.CrossEntropyLoss()
criterion2 = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

input_data = torch.randn(1, 10) # Input data of shape (batch_size, input_size)

# Forward pass
out1, out2 = model(input_data)

# Compute losses
loss1 = criterion1(out1, torch.randn(1, 3)) # MSE loss in first stream
loss2 = criterion2(out2, torch.tensor([1])) # Cross-entropy loss in second stream

# Compute total loss
total_loss = loss1 + loss2

# Backward pass
optimizer.zero_grad()
loss1.backward(create_graph = True)
optimizer.step()

# Visualize computational graph and gradients
make_dot((loss1), params=dict(model.named_parameters()))
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dbe0bc9eb97b488070f23028f850403e22ca82e.png" data-download-href="https://discuss.pytorch.org/uploads/default/6dbe0bc9eb97b488070f23028f850403e22ca82e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dbe0bc9eb97b488070f23028f850403e22ca82e_2_386x500.png" alt="image" data-base62-sha1="fEPdoxOCLWtskI0wR7R5GaiPVVQ" width="386" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dbe0bc9eb97b488070f23028f850403e22ca82e_2_386x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/d/6dbe0bc9eb97b488070f23028f850403e22ca82e_2_579x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/d/6dbe0bc9eb97b488070f23028f850403e22ca82e.png 2x" data-dominant-color="EDEFEF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">647×837 21.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ptrblck: <p>I don’t know if <code>torchviz</code> is able to display all related tensors and it also seems this library wasn’t updated in ~2 years, so I’m also unsure if it’s still being maintained or not.</p>

ShanakaRG: <p>yeah me too. because it gave the same graph for both loss1+loss2 and loss1*loss2. So do you know if there is a method to make sure that stream2 weights are updated using both loss1 and  loss2?. it is better if we can visualize the backpropagation path for a simple network like the above.<br>
I appreciate your consideration <a class="mention" href="/u/ptrblck">@ptrblck</a> !! thank you very much.</p>