Autograd row-wise of a tensor using PyTorch autograd and without for loop

Arka_Roy: <p>I am new in PyTorch and developing a deeponet code. I need to create a loss function where autograd is applied on each sample of model output. I used for loop, but the code is very slow. Is there any way to vectorize the for loop to get the grad for entire samples in one go. I am attaching the loss function code</p>
<blockquote>
<p>def derivative(dy: torch.Tensor, x: torch.Tensor, order: int = 1) → torch.Tensor:<br>
“”"<br>
This function calculates the derivative of the model at x_f<br>
“”"<br>
for i in range(order):<br>
dy = torch.autograd.grad(<br>
dy, x, grad_outputs = torch.ones_like(dy), create_graph=True, retain_graph=True<br>
)[0]<br>
return dy</p>
</blockquote>
<blockquote>
<p>def loss_phy_mod(model: nn.Module, y_data, x_data):<br>
yy=model(y_data,x_data)<br>
jacobian_rows = [derivative(yy[i,:],x_data,1)<br>
for i in range(y_data.shape[0])]<br>
jacobian = torch.stack(jacobian_rows)</p>
<pre><code>  loss=  torch.mean((y_data.flatten() - jacobian.flatten())**2)
  
  return loss
</code></pre>
</blockquote>

KFrank: <p>Hi Arka!</p>
<aside class="quote no-group" data-username="Arka_Roy" data-post="1" data-topic="206553" data-full="true">
<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/arka_roy/48/71528_2.png" class="avatar"> Arka_Roy:</div>
<blockquote>
<p>Is there any way to vectorize the for loop to get the grad for entire samples in one go.</p>
</blockquote>
</aside>
<p>Composing <a href="https://pytorch.org/docs/2.3/generated/torch.func.vmap.html" rel="noopener nofollow ugc">vmap</a> with <a href="https://pytorch.org/docs/2.3/generated/torch.func.grad.html" rel="noopener nofollow ugc">grad</a> will likely do what you want.</p>
<p>Best.</p>
<p>K. Frank</p>

Arka_Roy: <p>Thanks Frank, I tried using vmap. But still not getting</p>
<blockquote>
<p>import functorch</p>
<h1>Function to compute gradients of a single output element w.r.t input</h1>
<p>def compute_grad(y_pred_elem, x):<br>
grad, = torch.autograd.grad(y_pred_elem, x, create_graph=True, retain_graph=True)<br>
return grad</p>
<h1>Vectorize this function to compute gradients for all elements</h1>
<p>vectorized_grad = functorch.vmap(compute_grad, in_dims=(0, None))</p>
<h1>Compute the Jacobian</h1>
<p>jacobian = vectorized_grad(y_pred, x.view(1,100))</p>
</blockquote>
<p>getting a error element 0 of tensors does not require grad and does not have a grad_fn. but x having requires_grad True.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/arka_roy">@Arka_Roy</a> (and <a class="mention" href="/u/kfrank">@KFrank</a>),</p>
<p>The reason why it isn’t working with <code>torch.func.vmap</code> is that <code>torch.func.vmap</code> requires the entire process be within its ‘funtionalized’ approach, i.e. you can’t mix <code>torch.autograd</code> operations with <code>torch.func</code> when computing higher derivatives.</p>
<p>You can look at a previous answer I’ve shared on the forums, which focuses on using <code>torch.func</code> to compute the Hessian, here: <a href="https://discuss.pytorch.org/t/efficient-computation-of-hessian-with-respect-to-network-weights-using-autograd-grad-and-symmetry-of-hessian-matrix/156222/8" class="inline-onebox">Efficient computation of Hessian with respect to network weights using autograd.grad and symmetry of Hessian matrix - #8 by AlphaBetaGamma96</a></p>

Arka_Roy: <p>Dear <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>,<br>
Here  I am sharing a simple code for vectorized grad. One is the sequential approach and other one using vmap. First one working fine but second one not. I am not finding any error. Could you please try and correct it.</p>
<pre><code class="lang-auto">import torch

# Device agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)

def derivative(dy: torch.Tensor, x: torch.Tensor) -&gt; torch.Tensor:
    """
    This function calculates the derivative of the model at x_f
    """
    dy = torch.autograd.grad(
            dy, x, grad_outputs=torch.ones_like(dy), create_graph=True, retain_graph=True
                            )[0]
    return dy

def create_func(x: torch.tensor) -&gt; torch.tensor:
    yy = torch.empty(2, x.shape[0])
    y1 = x ** 2
    y2 = x ** 3
    yy[0, :] = y1
    yy[1, :] = y2
    return yy

x = torch.linspace(0, 1, 10).type(torch.float32).to(device)
x.requires_grad = True

yy = create_func(x).to(device)

# Sequential approach to calculate the Jacobian
jacobian_rows = [derivative(yy[i, :], x) for i in range(2)]
jacobian = torch.stack(jacobian_rows)

# Function to compute gradients of a single output element w.r.t input
def compute_grad(y_pred_elem, x):
    grad, = torch.autograd.grad(y_pred_elem, x, create_graph=True, retain_graph=True)[0]
    return grad

# Custom vectorization using torch.vmap
vectorized_grad = torch.vmap(compute_grad, in_dims=(0, None))
jacobian = vectorized_grad(yy, x)

print(jacobian)
</code></pre>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/arka_roy">@Arka_Roy</a>,</p>
<aside class="quote no-group" data-username="Arka_Roy" data-post="5" data-topic="206553">
<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/arka_roy/48/71528_2.png" class="avatar"> Arka_Roy:</div>
<blockquote>
<pre><code class="lang-auto"># Function to compute gradients of a single output element w.r.t input
def compute_grad(y_pred_elem, x):
    grad, = torch.autograd.grad(y_pred_elem, x, create_graph=True, retain_graph=True)[0]
    return grad

# Custom vectorization using torch.vmap
vectorized_grad = torch.vmap(compute_grad, in_dims=(0, None))
jacobian = vectorized_grad(yy, x)
</code></pre>
</blockquote>
</aside>
<p>So, you’re mixing <code>torch.autograd</code> with <code>torch.func</code> operations,</p>
<aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="4" data-topic="206553">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>The reason why it isn’t working with <code>torch.func.vmap</code> is that <code>torch.func.vmap</code> requires the entire process be within its ‘funtionalized’ approach, i.e. you can’t mix <code>torch.autograd</code> operations with <code>torch.func</code> when computing higher derivatives.</p>
</blockquote>
</aside>
<p>which leads to a zero gradient. You need a complete <code>torch.func</code> approach in order for <code>torch.func</code> to return the correct gradient. An example can be found below,</p>
<pre><code class="lang-auto">import torch
from torch.func import jacrev, vmap

device = 'cuda' if torch.cuda.is_available() else 'cpu'

def func(x):
  y1=x**2
  y2=x**3
  return torch.stack([y1,y2],dim=-1)

x = torch.linspace(0,1,10, dtype=torch.float32, device=device) #don't need requires_grad=True with torch.func

jacobian = vmap(jacrev(func, argnums=(0)), in_dims=(0))(x)
print(jacobian)
</code></pre>
<p>returns,</p>
<pre><code class="lang-auto">tensor([[0.0000, 0.0000],
        [0.2222, 0.0370],
        [0.4444, 0.1481],
        [0.6667, 0.3333],
        [0.8889, 0.5926],
        [1.1111, 0.9259],
        [1.3333, 1.3333],
        [1.5556, 1.8148],
        [1.7778, 2.3704],
        [2.0000, 3.0000]], device='cuda:0', grad_fn=&lt;ViewBackward0&gt;)
</code></pre>

Arka_Roy: <p>Dear <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>,<br>
Thanks a lot.</p>

Klae_zhou: <p>Hi, I have a same question, but I need to just get one layers gradient to save some time in my model. But the function ‘jacrev’ can’t let me to choose the exact layer’s parameter to backpropagate. <a href="https://discuss.pytorch.org/t/autograd-on-a-specific-layers-parameters/212541" class="inline-onebox">Autograd on a specific layer's parameters</a> This is my full problem.</p>

AlphaBetaGamma96: <p>You’ll need to share a minimal reproducible example of your issue (including the model) so I can understand what’s wrong.</p>

Klae_zhou: <p>Thank you for your answering. Here is my codes</p>
<pre><code class="lang-auto">act=torch.nn.Tanh()
model = nn.Sequential(
            nn.Linear(in_features=3, out_features=5),
            act,
            nn.Linear(in_features=5, out_features=1, bias=False)
        )
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

batch_size = 4
input_features = 3
input_data = torch.randn(batch_size, input_features, device=device)  #init input   
</code></pre>
<pre><code class="lang-auto">func_params=dict(model.named_parameters())
                
def fm(params,input):
    return functional_call(model,params,input.unsqueeze(0)).squeeze(0)

# loss function                   
def floss(func_params,input):      
    fx=fm(func_params,input)
    return fx

def grad(f,param):
    return torch.autograd.grad(f,param)
</code></pre>
<pre><code class="lang-auto">per_sample_grads =vmap(jacrev(floss,0), (None, 0))(func_params, input_data)

#i want to get the gradient of ['0.weight']
oweightgrads=vmap(grad(floss,func_params['0.weight']),(None,0))(func_params,input_data)
</code></pre>

AlphaBetaGamma96: <p>I’m pretty sure you need to pass all weights to the <code>vmap</code> call, then select the specific weight after the call. As your loss function depends on all the weights, you need to pass all weights to the <code>floss</code> function, even if you just want a subset of the weights.</p>
<p>So, something like this,</p>
<pre><code class="lang-auto">import torch 
from torch import nn
from torch.func import vmap, jacrev, functional_call

act=torch.nn.Tanh()
model = nn.Sequential(
            nn.Linear(in_features=3, out_features=5),
            act,
            nn.Linear(in_features=5, out_features=1, bias=False)
        )
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
params=dict(model.named_parameters())

batch_size = 4
input_features = 3
input_data = torch.randn(batch_size, input_features, device=device)  #init input   
                
def fm(params,input):
    return functional_call(model,params,input.unsqueeze(0)).squeeze(0)

def floss(params, input):
  return functional_call(model, params, input)

per_sample_grads = vmap(jacrev(floss, argnums=(0)), in_dims=(None,0))(params,input_data) 
print('per_sample_grads: ',per_sample_grads) # returns PyTree
oweightgrads = per_sample_grads['0.weight']
print('oweightgrads: ',oweightgrads) # returns our select weights (shape: oweightgrads:  torch.Size([4, 1, 5, 3]))
</code></pre>

Klae_zhou: <p>Thanks for your answer. So i can’t save time in this step by just computing some specific weights’ gradient. Is there some other approach? Because I just need one layers gradient in my algorithm.</p>

AlphaBetaGamma96: <p>But your loss depends on the output of the network, which depends on <strong>all</strong> the weights. I don’t see how you can calculate the gradient of the loss for a specific part of the network without using all the weights?</p>

Klae_zhou: <p>Yeah, but i found this approach works.</p>
<pre><code class="lang-auto">gradient=grad(floss(func_params,input_data)[0],func_params['0.weight'])
</code></pre>
<p>However, it just works when I choose one batch. Which means I can just use it with for loop.</p>

AlphaBetaGamma96: <p>If this the gradients of the loss (as in averaged over all samples), then taking the gradient with respect to 1 particular weight? Why not just <code>vmap</code> that call?</p>

Klae_zhou: <p>Sorry, I may not catch the point. Do your mean call vmap like this ?</p>
<pre><code class="lang-auto">oweightgrads=vmap(grad(floss,func_params['0.weight']),(None,0))(func_params,input_data)
</code></pre>

Klae_zhou: <p>Thank you so much for taking the time to help my question. I’m happy to share that I haved solved it with the following approach. Here is the code:</p>
<pre><code class="lang-auto">c=func_params['0.weight']
print(c.shape)
basis_vector=torch.eye(batch_size).to(device)
basis_vector=basis_vector.unsqueeze(-1)
loss=floss(func_params,input_data)
def grad(v):
    return torch.autograd.grad(loss,c,v,retain_graph=True)

g=vmap(grad)(basis_vector)
</code></pre>
<p>But when run this code, there has a warning</p>
<pre><code class="lang-auto">UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at ../aten/src/ATen/cuda/CublasHandlePool.cpp:135.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
</code></pre>