How to efficiently compute gradient for each training sample?

aaronhan223: <p>Hi folks,</p>
<p>There is a problem that has bothered me for quite a long time. Assume we are minimizing a loss function <img src="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b6588a1ec454fb4d54dda0feece18ab0fc8d1033.png" alt="WechatIMG4" data-base62-sha1="q16AS0fQosmNy9WsikYSuBqJFD5" width="148" height="62"> parameterized by <img src="https://discuss.pytorch.org/uploads/default/original/3X/f/8/f88c930b275fde4fcbba22ef140a1aa30971e5b5.png" alt="WechatIMG3" data-base62-sha1="zsLAHiSGOr5n0lweXzOWXcoGJLL" width="36" height="52">, on samples <img src="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c76d2a8c164a9aa68c5181c86324b3f5aa57723c.png" alt="WechatIMG5" data-base62-sha1="sscU3s8Bb7YMsHjKQnZ88SH33dO" width="240" height="72"> using SGD, where M is the mini-batch size. Since the PyTorch autograd can only be implicitly created for <strong>scalar</strong> outputs, I am wondering if there is any efficient way to compute the gradient for <strong>each sample</strong>, i.e., <img src="https://discuss.pytorch.org/uploads/default/original/3X/4/4/44bfc8bf58330e523bdeb64b0ccc1cbfedb23903.png" alt="image" data-base62-sha1="9ObmxYTWiu2ykGeZCgh28hhMKIP" width="254" height="76">, without setting the batch size equals to 1 and compute in a for loop (which is too slow)?</p>
<p>Thank you for your help!</p>

Sunshine352: <p>loss.backward(gradient=torch.ones_like(loss))<br>
The shape of loss is (B, )，which is the batch size</p>

aaronhan223: <p>Hi Zhaomang,</p>
<p>Thanks so much for your reply! If the loss is set to be a non-scaler, i.e., the sample loss:</p>
<pre><code class="lang-auto">loss = F.cross_entropy(predictions, results, reduction='none')
loss.backward(gradient=torch.ones_like(loss))
</code></pre>
<p>I’m still confused about how to get the gradient on each sample, if I run the following code:</p>
<pre><code class="lang-auto">for params in model.parameters(): 
    print(params.grad)
</code></pre>
<p>I expect to get <code>B</code> times the number of tensors returned in the mini-batch case, since there should be an individual result for each sample.</p>
<p>In addition, it would be very helpful if you could point me to the document of Tensor.backward() function (if any), so that I can learn more on this. Thanks again!</p>

Yaroslav_Bulatov: <p>If you just have conv/linear layers, you could use this – <a href="https://github.com/cybertronai/autograd-hacks#per-example-gradients" rel="nofollow noopener">https://github.com/cybertronai/autograd-hacks#per-example-gradients</a></p>

Sunshine352: <p>For the params.grad, it is already the accumulated gradient of B samples.<br>
But if you need to obtain the gradient on each sample, you can set<br>
X.retain_grad() before backward. And then through X.grad to get the gradient of loss on each sample.</p>

konstmish: <p>Are you going to extend this to other popular layer types, e.g. ReLu and batchnorm? This would make it useful for most widely used architectures such as Resnets</p>

Yaroslav_Bulatov: <p>It should work for those architectures as well. What’s missing is support for other layers with trainable parameters, (ie the multiheadattention layer)</p>

konstmish: <p>Not with batch norm. If I understood your implementation correctly, layers with batch normalization are just skipped:</p>
<pre><code class="lang-auto">_supported_layers = ['Linear', 'Conv2d']

if layer_type not in _supported_layers:
    continue
</code></pre>
<p>Which means that for batch normalization layer (which is trainable), grad1 will not be created.</p>

Yaroslav_Bulatov: <p>Ah right, the gamma, beta parameters. It seems feasible to extend computation explained here to per-example computation  <a href="https://kratzert.github.io/2016/02/12/understanding-the-gradient-flow-through-the-batch-normalization-layer.html" rel="nofollow noopener">https://kratzert.github.io/2016/02/12/understanding-the-gradient-flow-through-the-batch-normalization-layer.html</a></p>

aaronhan223: <p>Can the current version of autograd-hacks still be used if I apply other activation functions such as Sigmoid / tanh / ReLU? Thanks.</p>

Yaroslav_Bulatov: <p>Yes, any activation function that works with PyTorch autograd should also work here</p>

offsongjh: <p>Thanks for per-example-gradients function!</p>
<p>Can I use this function by only importing torch?</p>
<p>I mean will per-example-gradients be an official function of Pytorch in the future?</p>

dampeyrou_charles: <p>Hi everyone !</p>
<p>Thank you Yaroslav for your github !</p>
<p>I have found a recent library called functorch that does compute per-sample gradients, check this out if you have similar problems : <a href="https://pytorch.org/functorch/stable/notebooks/per_sample_grads.html" class="inline-onebox" rel="noopener nofollow ugc">Per-sample-gradients — functorch 0.1.1 documentation</a></p>