Simple use case: Compete per sample gradient with autograd

Tao_Wang1: <p>Hi, I came across a simple use case of pytorch autograd function, I have batch of input to a network X of shape [N, H], and the loss is L of shape [N,]. now I want to get the graident of each loss sample independently w.r.t to the network layer weight w, my workaround now is to loop over the N samples, and call the torch.autograd.grad function multiple times:</p>
<pre><code class="lang-auto">def compute_grad(loss_batch, weight):
"""
loss_batch, batched loss for multiple samples, [N,]
weight, the network layer weight for computing the gradients
"""
    grads = []
    for i in range(len(loss_bbox)):
        grad_persample = torch.autograd.grad(loss_batch[i], weight, retain_graph=True)[0]
        grads.append(grad_persample)
    grads = torch.stack(grads)
    return grads
</code></pre>
<p>is there a batched way to implement this simple use case?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/tao_wang1">@Tao_Wang1</a>,</p>
<p>If you want to batch (or vectorize) over a function, you’ll want to have a look at the <code>torch.func</code> namespace. The documentation is here: <a href="https://pytorch.org/docs/stable/func.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func — PyTorch 2.4 documentation</a></p>

Tao_Wang1: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> thanks for your suggestion, can torch.func parallelize the gradient computation? if so, could you pls suggest simple implementation code. As currently, the major problem with my naive loop implementation is that it is very slow.</p>

AlphaBetaGamma96: <p>It can vectorize over a for-loop, yes. Just bare in mind, you can’t mix <code>torch.autograd.grad</code> with <code>torch.func</code>, so you’ll need to re-write how your loss function is defined.</p>

Tao_Wang1: <p>I’m sorry, do you mean it actually can parallelize the torch.autograd.grad computation? and I need to implement the loss function? but currently the loss is already computed from the network, and the gradient is w.r.t to a single layer in the network</p>

AlphaBetaGamma96: <p>The way <code>torch.func</code> works (at least to my understanding) is that it defines a function for a single sample (i.e. computing a derivative) and then you compose this with <code>torch.func.vmap</code> to vectorize over an entire batch of samples. So, all samples are computing in parallel instead of using a for-loop, like you might with <code>torch.autograd.grad</code>.</p>
<p>So if you have function that computes the derivative of the loss with respect to some inputs, <code>torch.func</code> can vectorize over this function for multiple inputs (alleviating the overhead of calling <code>autograd</code> multiple times instead of just once).</p>
<p>Edit: I’d also recommend looking through the jacobian/hessian documentation to see how <code>vmap</code> composes with derivative calculations too. Documentation here: <a href="https://pytorch.org/functorch/stable/notebooks/jacobians_hessians.html" class="inline-onebox" rel="noopener nofollow ugc">Jacobians, Hessians, hvp, vhp, and more: composing functorch transforms — functorch nightly documentation</a></p>

Tao_Wang1: <p>Thanks for the suggestion, the current difficulty is that the function is implicitly computed in the my model forward pass (not like the example here, the function is simplely defined with a few lines), and I just want mimic torch.autograd.grad to pass in the loss and weight, then the computation graph could be used to compute gradient</p>

AlphaBetaGamma96: <p>If you’re working with an <code>nn.Module</code> object, you need to use <code>torch.func.functional_call</code> to ‘functionalize’ your model. Here’s the docs: <a href="https://pytorch.org/docs/stable/func.migrating.html#nn-module-utilities" class="inline-onebox" rel="noopener nofollow ugc">Migrating from functorch to torch.func — PyTorch 2.4 documentation</a></p>

Tao_Wang1: <p>thanks for the suggestion, I found the per-sample graident example at <a href="https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html" class="inline-onebox" rel="noopener nofollow ugc">Per-sample-gradients — functorch nightly documentation</a>, it seems that the example ‘functionalize’ a model, but currently I just want compute the gradient w.r.t to a single layer, and implement this part of computation within the module’s forward function, I’m wondering if that is possible?</p>

AlphaBetaGamma96: <p>I think you can just pass the key for that layer when you compute the gradients, but I’m not 100% sure.</p>

Tao_Wang1: <p>I’ll have a try, thank you. But I really think it could be good if torch.autograd.grad could incorporate a per-sample gradient computation feature, it currently have <em>is_grads_batched</em>, but seems not for this use case?</p>

AlphaBetaGamma96: <p>Have you modified the <code>grad_outputs</code> arg of <code>torch.autograd.grad</code>? Like so,</p>
<pre><code class="lang-auto">grad_persample = torch.autograd.grad(loss_batch, weight, torch.ones_like(loss_batch), retain_graph=True)[0]
</code></pre>

Tao_Wang1: <p>yes, but that returns the gradient with the same shape as <code>grad_persample = torch.autograd.grad(loss_batch[i], weight, retain_graph=True)[0]</code>, e.g., [4800, 1024] (which is the shape of the weight), the value is different though. I would expect per-sample gradient as shape, e.g., [N, 4800, 1024], where N is the number of samples.</p>

AlphaBetaGamma96: <p>Yes, that might implicitly sum, but I’m not too sure at the moment. It can be used to compute per-sample gradients, but it can become quite limited. That’s why I recommended <code>torch.func</code> as it’s much more efficient, but it does require some re-writes of your code away from the <code>torch.autograd</code> namespace.</p>

Tao_Wang1: <p>I understand, thanks for the suggestion. Is it possible to make a new feature request to the pytorch library to add such functionality?</p>

AlphaBetaGamma96: <p>You might be able to do it with <code>torch.autograd.grad</code>, it’d be best to get a dev’s opinion (<a class="mention" href="/u/ptrblck">@ptrblck</a>, apologizes for the tag!)</p>
<p>For new feature request, it’s best your make a request on their github repo.</p>

Tao_Wang1: <p>I see, thanks for the suggestion and I have made a new feature request on pytorch github repo, <a class="mention" href="/u/ptrblck">@ptrblck</a> could you please advise on this topic?</p>

DanChigo: <p>Hi! Do you have a link to the feature request you made, and are their any updates on this, I would love to know!</p>