Understanding torch auto-grad

Innat: <p>Hello.<br>
I’m trying to understand the following torch code.</p>
<pre><code class="lang-python">ran_model = Random()
ran_optim = torch.optim.SGD(
    ran_model.parameters(),0.01
)

model = Model()
optim = torch.optim.SGD(
    model.parameters(),0.01
)
model_params = list(model.parameters())
</code></pre>
<pre><code class="lang-python">for i in range(5):
    loss_mod  = torch.mean(torch.log(model.forward(x)[:,i]))
    loss_rand = torch.mean(torch.log(model.forward(y)[:,i]))
    
    model_grad = torch.autograd.grad(loss_mod, model_params)
    rand_grad  = torch.autograd.grad(
        loss_rand, model_params, create_graph=True
    )
    
    loss = 0
    for j in range(len(model_grad)):
        a = model_grad[j]
        b = rand_grad[j]
        loss = loss + torch.stack(
            [a,b], dim=0
        ).sum(dim=0).sum(dim=0).mean(0)
        
    ran_model.zero_grad()
    loss.backward()
    ran_optim.step()
    
    new_loss = -torch.mean(torch.log(model.forward(y))) 
    optim.zero_grad()
    new_loss.backward()
    optim.step()
    print('loss value  ', loss, new_loss)
</code></pre>
<hr>
<p>Here,</p>
<ol>
<li>Get losses after forward pass, <code>loss_mod</code> and <code>loss_rand</code>.</li>
</ol>
<pre><code class="lang-auto">loss_mod  = torch.mean(torch.log(model.forward(x)[:,i]))
loss_rand = torch.mean(torch.log(model.forward(y)[:,i]))
</code></pre>
<ol start="2">
<li>Compute grads of <code>loss_mod</code> and <code>loss_rand</code> w.r.t  <code>list(model.parameters())</code>.</li>
</ol>
<pre><code class="lang-auto">model_grad = torch.autograd.grad(loss_mod, model_params)
rand_grad  = torch.autograd.grad(
  loss_rand, model_params, create_graph=True
)
</code></pre>
<ol start="3">
<li>Update the above grads, just modify them in some ways.</li>
</ol>
<pre><code class="lang-auto">loss = some_method(model_grad, rand_grad)
</code></pre>
<ol start="4">
<li>Make grad zero to the <code>ran_model = Random()</code> model. And compute the gradients from the above-updated</li>
</ol>
<p>loss (comes from two gradients) w.r.t <code>ran_model</code> model’s params. And next, update the state of <code>rand_model</code> with its optimizer.</p>
<pre><code class="lang-auto">ran_model.zero_grad()
loss.backward()
ran_optim.step()
</code></pre>
<p>I’m confused about step 4.<br>
How the gradients are calculated here (<code>loss.backward()</code>) w.r.t <code>ran_model</code>? The loss is not achieved from <code>ran_model</code>. So, there should be no connection. How does Torch handle these cases? These loss values which come from two gradients (<code>'model_grad</code> and <code>rand_grad</code>) are achieved from <code>model = Model()</code> and not from <code>ran_model = Random()</code>.</p>

soulitzer: <p>If there is no connection, then gradients wouldn’t be calculated wrt ran_model. Unlike .grad(), .backward() doesn’t complain if there is no connection, it would just accumulate the gradients to the inputs that are connected in the graph.</p>

Innat: <p>Could you please elaborate?</p>
<blockquote>
<p>It would just accumulate the gradients to the inputs that are connected in the graph.</p>
</blockquote>
<p>In the code, while computing the grad of <code>loss_rand</code> w.r.t. <code>model_params</code>, we set <code>create_graph=True</code>. I think that means, here somehow <code>loss_rand</code> and <code>model_params</code> getting connected, thus allowing us to compute grad of grads in the later part.</p>
<pre data-code-wrap="python"><code class="lang-python">rand_grad  = torch.autograd.grad(
  loss_rand, model_params, create_graph=True
)
</code></pre>

soulitzer: <p>Yes that means params.grad is now a function of loss_rand potentially, hence connected</p>
<blockquote>
<p>It would just accumulate the gradients to the inputs that are connected in the graph.</p>
</blockquote>
<p>Ah, all I mean here is that the preferred way of figuring out what is connected to what is with .grad() rather than .backward()</p>

Innat: <p>Thank you.<br>
Could you please suggest any blogs/docs that describe the behavior of such cases of grad calculation in the torch?</p>

soulitzer: <p>If you are trying to understand higher order graphs, e.g. double backward, there’s this one that may be helpful <a href="https://pytorch.org/tutorials/intermediate/custom_function_double_backward_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Double Backward with Custom Functions — PyTorch Tutorials 2.0.1+cu117 documentation</a>, technically it is describing how one would implement custom functions that can support autograd, but it has some visualizations of what the graph would look like</p>