Speeding up gradient computation; instead of using a for loop

username: rajoy99
cooked: <p>I have a torch tensor named HSNR of size 4096. I am calculating the gradient of it w.r.t. NN parameters in this approach. However the process is extremely slow. How can I speed it up ?</p>
<pre><code class="lang-auto">for i in range(4096):

     first_gradient_list[i] = torch.autograd.grad(HSNR[i], NN.parameters(),retain_graph=True). 
</code></pre>

username: Grant_Norman
cooked: <p>It might be hard to answer this how you were hoping without more information, but I’d suggest looking into the <code>functorch</code> documentation for this. It’s a semi-recent addition as <code>torch.func</code> that enables faster and more math-like differentiation.</p>
<p>You could try here for the actual documentation:<br>
<a href="https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html#per-sample-grads-the-efficient-way-using-functorch" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html#per-sample-grads-the-efficient-way-using-functorch</a></p>
<p>I created a function like this, where <code>b</code> is batch dimension, <code>d</code> is dimension of the input to the forward call, and <code>n</code> is the number of parameters.</p>
<pre><code class="lang-auto">    def dn_dtheta(self, x):
        """(b, d) -&gt; (b, n)"""
        dn_dtheta_dict = torch.func.jacrev(self._forward_w_params)(self.get_theta(), x)
        dn_dtheta = dict_to_vec_fl(dn_dtheta_dict)
        return dn_dtheta
</code></pre>
<p>in my <code>nn.Module</code>, where I use</p>
<pre><code class="lang-auto">    def _forward_w_params(self, theta, x):
        """(n, d) -&gt; (b)"""
        return torch.func.functional_call(model, theta, x)
</code></pre>
<p>although my forward function returns just a scalar for each batch <code>(b,d) -&gt; (b)</code>.</p>

username: rajoy99
cooked: <p>Thanks.</p>
<p>The problem domain is : solving burgers equations with Physics Informed Neural Networks.</p>
<p>Following the documentation I am trying to solve the problem like this. However, Getting the error : “RuntimeError: One of the differentiated Tensors does not require grad”</p>
<pre><code class="lang-auto">    data = torch.hstack((X,T))
    # Define a function to compute the loss for a single input 
    def compute_loss_stateless_model(params, buffers, sample):

        batch = sample.unsqueeze(0)

        u_pred = fmodel(params, buffers, batch) 

        # print(sample)
        print("Samples Req Grad : ",sample.requires_grad)
        
        print(batch.requires_grad)

        u_t = torch.autograd.grad(
            u_pred, batch[0][1],
            grad_outputs=torch.ones_like(u_pred),
            retain_graph=True,
            create_graph=True)[0]
    
        u_x = torch.autograd.grad(
            u_pred, batch[0][0], 
            grad_outputs=torch.ones_like(u_pred),
            retain_graph=True,
            create_graph=True)[0]
        
        u_xx = torch.autograd.grad(
            u_x, batch[0][0], 
            grad_outputs=torch.ones_like(u_x),
            retain_graph=True,
            create_graph=True)[0]

        residual = u_t + (u_pred * u_x) - (nu * u_xx)

        # Calculating the Half of Squared Euclidean Norm : half_squaredENorm_residual HSNR
        HSNR = 0.5 * (residual ** 2)
        
        return HSNR
    
    ft_compute_grad = grad(compute_loss_stateless_model) 

    ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, 0, 0))

    ft_per_sample_grads = ft_compute_sample_grad(params, buffers, data)
</code></pre>
<p>To debug I found : inside the compute_loss_stateless_model function the requires_grad flag of the sample becomes False. Is there any way around ?</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/rajoy99">@rajoy99</a>,</p>
<p>Expanding on <a class="mention" href="/u/grant_norman">@Grant_Norman</a>’s answer, you should ideally only use <code>torch.func</code> functions when computing per-sample gradients and not mix <code>torch.func</code> calls with <code>torch.autograd</code> calls (at least not for the calculation of the per-sample gradient), that’s why you get a <code>.grad</code> not required error.</p>
<p>Within the <code>torch.func</code> namespace when you vectorize over a function, it will no longer see the batch dimension and effectively redefines your function to work for a ‘single’ sample and then computes all these samples in parallel, so you need to remove your indexing on the input Tensor as well (<code>torch.func</code> can’t see the batch dim!)</p>
<p>Also, as you have noticed the arguments passed to <code>torch.func</code> calls may have their <code>.grad</code> flags disabled as it returns the gradients directly rather than populating the <code>.grad</code> attribute.</p>
<p>To see all the changes between <code>functorch</code> and <code>torch.func</code>, have a look at the migrating from <code>functorch</code> to <code>torch.func</code> docs <a href="https://pytorch.org/docs/stable/func.migrating.html" rel="noopener nofollow ugc">here</a>.</p>
<p>Your code above should be re-written to something like this,</p>
<pre><code class="lang-auto">net = Model(*args, **kwargs) #nn.Module

params = dict(net.named_parameters()) #params are now dicts in torch.func (rather than lists in functorch)

from torch.func import functional_call, vmap, jacrev, grad 

def fcall(params, sample): #functionalized call (with params as input)
  return functional_call(net, params, sample)

def compute_loss(params, sample):
  u_pred = fcall(params, sample)

  u_grads = grad(fcall, argnums=(1))(params, sample) #assume model outputs scalar
  #u_t and u_x can be spilt separately (depending on their shape and if they're independent of each other)

  u_xx =  jacrev(jacrev(fcall, argnums=(1)), argnums=(1))(params, sample) #reverse-over-reverse Hessian calc. 
  
  residual = u_t + (u_pred*u_x) - (nu*u_xx) #sum (be careful of the shapes here, you have a matrix and vector and try to avoid broadcasting issues).
  
  HSNR = 0.5 * (residual**2)
  return HSNR

losses = vmap(compute_loss, in_dims=(None,0))(params, samples)
</code></pre>
<p>Now granted, the example above is for a model with only <code>params</code> and not <code>buffers</code>, but changing that should be minimal (I just don’t know how your model is defined, but something like the <code>torch.func.stack_module_state</code> can give an example of how to combine <code>buffers</code> with <code>params</code> when <code>vmap</code>-ing.</p>
<p>You’ll then need to modify the <code>grad</code> and <code>jacrev</code> calls to include the <code>buffers</code> object (but that won’t be too difficult).</p>

username: rajoy99
cooked: <p>Thank you. I made some modifications and it solved my problem of per sample gradient computation. I have verified the computations are correct by comparing with my previous brute force solution.</p>
<p>Still, I have a question. Since I am computing the gradient wrt NN parameters, I get the result as a dictionary of tensors.</p>
<pre><code class="lang-auto">HSNR_grad.keys()
&gt;&gt;&gt; dict_keys(['0.weight', '0.bias', '2.weight'.....]) 
</code></pre>
<pre><code class="lang-auto">HSNR_grad['2.weight'].shape 
&gt;&gt;&gt; torch.Size([4096, 20, 20])
</code></pre>
<p>From which I construct the list_of_gradients by running a for loop over the dictionary.  Like this :</p>
<pre><code class="lang-auto">
   keylist = HSNR_grad.keys()

   first_gradient_list = []

   for i in range(4096):
       gradients_at_index =[]

       for k in keylist:
           gradients_at_index.append(HSNR_grad[k][i])
       first_gradient_list.append(gradients_at_index)

</code></pre>
<p>Is there any approach which will allow me to compute the gradient in a per sample list format (where each entry is the gradient for that particular sample). It will be more efficient than running the for loop again.</p>

username: AlphaBetaGamma96
cooked: <p>You could try using <code>torch.unbind</code> to convert the batch dim to a list, docs <a href="https://pytorch.org/docs/stable/generated/torch.unbind.html" rel="noopener nofollow ugc">here</a></p>