How to get the gradients of each sample in a batch efficiently

likhilb9: <p>Can anyone help me with how to get the gradients for each sample in a mini-batch efficiently, not the one in the original forum.</p>
<p>Here is my code.<br>
criterion = nn.CrossEntropyLoss(reduction=‘none’)<br>
loss_fk = criterion(model(inputs), labels)<br>
[torch.autograd.grad(loss, model.parameters(), retain_graph=True) for loss in loss_fk]</p>
<p>Thanks in advance.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/likhilb9">@likhilb9</a>,</p>
<p>You can easily do this within the <code>torch.func</code> framework as follows,</p>
<pre><code class="lang-auto">from torch.func import vmap, grad, functional_call

criterion = nn.CrossEntropyLoss(reduction=‘none’)

def fcall(params, inputs):
  return functional_call(model, params, inputs) #functional call (params now  input).

def loss_fn(params, inputs, labels):
  outputs = fcall(params, inputs)
  loss = criterion(outputs, labels)
  return loss

params = dict(model.named_parameters()) #params must be dict (not list)

loss_wrt_params = vmap(grad(loss_fn, argnums=(0)), in_dims=(None, 0, 0))(params, inputs, labels) #torch.func.grad for reference.
</code></pre>
<p>Also, make sure to wrap code between three backticks ``` to get it formatted properly</p>
<p>EDIT: Corrected missing <code>vmap</code> call.</p>

likhilb9: <p>The above code is not working properly, and there is some issues with batch norm, and also I could’nt see you using vmap</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/likhilb9">@likhilb9</a>,</p>
<p>I’ve made the correction for the <code>vmap</code>/<code>grad</code> call.</p>
<p>In order to deal with batch norm layers, you can use the <code>torch.func.replace_all_batch_norm_modules_</code> function, docs <a href="https://pytorch.org/docs/stable/generated/torch.func.replace_all_batch_norm_modules_.html#torch.func.replace_all_batch_norm_modules_" rel="noopener nofollow ugc">here</a>.</p>
<p>You can find more information on patching the batch norm layer in the documentation here: <a href="https://pytorch.org/docs/stable/func.batch_norm.html" class="inline-onebox" rel="noopener nofollow ugc">Patching Batch Norm — PyTorch 2.2 documentation</a></p>

likhilb9: <p>mat1 and mat2 shapes cannot be multiplied (1024x25 and 400x120)</p>
<p>The same code is giving this error. Only while using VMap.<br>
I have removed vmap and just tried implement the loss_fn function. It returned me the 64 losses.<br>
The error only pops when I use this line vmap(grad(loss_fn, argnums=(0)), in_dims=(None, 0, 0))(params, inputs, labels) <span class="hashtag-raw">#torch</span>.func.grad for reference.</p>
<p>Here is my test model.<br>
class SimpleCNN(nn.Module):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()<br>
self.conv1 = nn.Conv2d(3, 6, 5)<br>
self.pool = nn.MaxPool2d(2, 2)<br>
self.conv2 = nn.Conv2d(6, 16, 5)<br>
self.fc1 = nn.Linear(16 * 5 * 5, 120)<br>
self.fc2 = nn.Linear(120, 84)<br>
self.fc3 = nn.Linear(84, 10)</p>
<p>def forward(self, x):<br>
x = self.pool(F.relu(self.conv1(x)))<br>
x = self.pool(F.relu(self.conv2(x)))<br>
x = torch.flatten(x, 1)<br>
x = F.relu(self.fc1(x))<br>
x = F.relu(self.fc2(x))<br>
x = self.fc3(x)<br>
return x</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="likhilb9" data-post="5" data-topic="197662">
<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/likhilb9/48/67935_2.png" class="avatar"> likhilb9:</div>
<blockquote>
<p>x = torch.flatten(x, 1)</p>
</blockquote>
</aside>
<p><code>vmap</code> removes the batch dimension, so you’ll need to use relative dims within <code>torch</code> functions. For example, a <code>tensor</code> of  shape <code>[B, N]</code> would be viewed as shape <code>[N]</code> by <code>vmap</code> within your <code>torch.flatten(x,1)</code> function.</p>

likhilb9: <p>Any other solutions please?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/likhilb9">@likhilb9</a>,</p>
<aside class="quote no-group" data-username="likhilb9" data-post="5" data-topic="197662">
<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/likhilb9/48/67935_2.png" class="avatar"> likhilb9:</div>
<blockquote>
<p>x = torch.flatten(x, 1)</p>
</blockquote>
</aside>
<p>The output size of the intermediate vector is of a different size than the input nodes in the <code>self.fc1</code> layer as stated in the error message. You can fix this by changing the number of inputs nodes in <code>self.fc1</code>.</p>
<p>You also need to change the <code>torch.flatten</code> <code>dim</code> arg to be relative rather than absolute, as the batch-size is removed when <code>vmap</code>ping, an example code can be found below.</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.func import vmap, jacrev, functional_call
import torch.nn.functional as F

class SimpleCNN(nn.Module):

  def __init__(self):
    super(SimpleCNN, self).__init__()
    
    self.conv1 = nn.Conv2d(3, 6, 5)
    self.pool = nn.MaxPool2d(2, 2)
    self.conv2 = nn.Conv2d(6, 16, 5)
    self.fc1 = nn.Linear(64, 120) #&lt;&lt;&lt; wrong amount of input nodes (changed to 64)
    self.fc2 = nn.Linear(120, 84)
    self.fc3 = nn.Linear(84, 10)

  def forward(self, x):
    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    x = torch.flatten(x, -3) #changed from 1 to -3 (relative dims, aka 3rd from the right rather than 1st from left)
    x = F.relu(self.fc1(x))
    x = F.relu(self.fc2(x))
    x = self.fc3(x)
    return x
    
net = SimpleCNN() #the network

def fcall(params, x):
  return functional_call(net, params, x)
  
params = dict(net.named_parameters())
x = torch.randn(4, 3, 20, 20) #random input size

y = net(x)

grad_params = vmap(jacrev(fcall, argnums=(0)), in_dims=(None,0))(params, x)
print(grad_params) #returns per-sample gradients with respect to params
</code></pre>