Quickly get individual gradients (not sum of gradients) of all network outputs

Quan_Vuong: <p>Right now, I’m doing:</p>
<pre><code>output_gradients = []

for output in net_outputs:
    tmp_grad = {}
    net.zero_grad()
    output.backward(retain_graph=True)
    for name, param in net.named_parameters():
         tmp_grad[name] = param.grad
    output_gradients.append(tmp_grad)
</code></pre>
<p>Since I have to call backward on each output, the backward passes are not parallelized and thus, the code is pretty slow.</p>
<p>Is there a faster way? Thanks!</p>

smth: <p>by default we only support accumulated gradients, so this is not easy to do.<br>
If you dont have memory constraints, you can use the <code>torch.autograd.grad</code> interface to compute separate gradients in one shot. (they wont be in <code>.grad</code>, but will be explicitly returned, so some manual book-keeping is needed)</p>

Quan_Vuong: <p>Thanks for the tips! Though we were able to rewrite the math to avoid doing this.</p>

fartash: <p><a class="mention" href="/u/smth">@smth</a> Can you elaborate more on how to use the <code>torch.autograd.grad</code> interface to compute separate gradients in one shot? Thanks.</p>

evcu: <p><code>grads = torch.autograd.grad(loss,parameters,retain_graph=True)</code><br>
would return the gradients as a tuple matching exactly the parameters provided. So that you don’t need to make a for loop to get the gradients as the code example above. Same size same order as the parameters provided.</p>

mercurialLMR: <p>Is it possible to do backward propagation in parallel way? Thanks!</p>

jusjusjus: <p>Dear <a class="mention" href="/u/evcu">@evcu</a>, in your example, <code>loss</code> is scalar and <code>parameters</code> is a vector.  However, the question is for multiple losses, and computing these without summation.  For me, <a class="mention" href="/u/smth">@smth</a>’s reference to <code>torch.autograd.grad</code> wasn’t helpful either.  Its docs start with: “Computes and returns the sum of gradients of outputs w.r.t. the inputs”, and that seems to be what it does.  Under what configuration does <code>grad</code> return individual gradients (at the expense of memory)?</p>
<p>As aside, This is a pretty useful feature for differential privacy.</p>

Skinish: <p>Exactly! Any updates on this? Would love to have an option to get individual gradients.</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/skinish">@Skinish</a>,</p>
<p>Look at the functorch library, which is designed to calculate per-sample gradients efficiently. Their github repo can be found <a href="https://github.com/pytorch/functorch/" rel="noopener nofollow ugc">here</a></p>

Skinish: <p>For what I see in <a href="https://github.com/pytorch/functorch/#working-with-nn-modules-make_functional-and-friends" rel="noopener nofollow ugc">here</a>, you need to perform the forward call in parallel to calculate the gradients. However, that consumes too much memory (RAM or GPU memory).<br>
What if we already have a single (X, Y) and just want to calculate gradients for all the elements in Y?</p>

AlphaBetaGamma96: <p>If you’re hitting an OOM error, you can always ‘chunk’ the vmap operation. So send multiple mini-batches and then just concatenate the results.</p>
<p>In the example you shared, a ‘chunk’ version would be something like this,</p>
<pre><code class="lang-auto">import torch
from functorch import make_functional, vmap, grad

model = torch.nn.Linear(3, 3)
data = torch.randn(64, 3)
targets = torch.randn(64, 3)

func_model, params = make_functional(model)

def compute_loss(params, data, targets):
    preds = func_model(params, data)
    return torch.mean((preds - targets) ** 2)

per_sample_grads = vmap(grad(compute_loss), (None, 0, 0))(params, data, targets)

out_full = vmap(grad(compute_loss), (None, 0, 0))(params, data, targets)

nchunks=8                   #number of chunks
nparams=len(params) #number of params, 2 in this case (weight &amp; bias)

#perform vmap in chunks 
out = [vmap(grad(compute_loss), in_dims=(None, 0, 0))(params, data_chunk, targets_chunk) for data_chunk, targets_chunk in zip(data.chunk(nchunks), targets.chunk(nchunks))] 
#recursively flatten list
out = [item for sublist in out for item in sublist] 
#re-map to correct shape and concatenate
out_chunk = [torch.cat([out[nparams*chunk+p] for chunk in range(nchunks)], dim=0) for p in range(nparams)] 

#check both methods for completeness
for i in range(2):
  print("Param: %i Match? %s" % (i, torch.allclose(out_full[i], out_chunk[i])))
"""
Returns
Param: 0 Match? True
Param: 1 Match? True
"""
</code></pre>
<p>If you still have an issue, you can open an issue on functorch’s github repo <a href="https://github.com/pytorch/functorch/issues/" rel="noopener nofollow ugc">here</a>. There is some development of to ‘chunk’ vmap as seen in issue <a href="https://github.com/pytorch/functorch/issues/680" rel="noopener nofollow ugc">#680</a> but it’s still under development I believe!</p>

Skinish: <p>Thanks for the reply <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
I opened an <a href="https://github.com/pytorch/functorch/issues/1010" rel="noopener nofollow ugc">issue</a> as you suggested</p>