Is there any way to calculate the per sample gradients with batch norm.
I guess Opacus convert the batch norm to Group norm and then calculate the per sample gradients is there any efficient way to calculate?
from nn.utils._per_sample_grad import call_for_per_sample_grads
res = call_for_per_sample_grads(model)(input)
loss(res, labels).backward()
for param in model.parameters():
param.grad_sample # instead of the grad, per sample gradients are in the grad_sample field
The above code does not work for batch norm layers