soulitzer:

I actually don’t quite understand what you are going for yet, but if you wanted to alter the gradients of both nn.Linear layers you’d need to individually register hooks for four tensors. Or do “for each mod.parameters()” twice.


I’ll make it a little bit clearer about that. I see what you’ve done, and it’s ok if you have two modules, but in my case, as in the image, it is the same module, thus the same parameters, so if I register a hook on them, it will be called actually twice in the backward path of the one called ‘am_loss’, and what I seek to achieve is a gradient divided by 10 only on the first time of two where the backward of the parameter is called.
So if I have only one module and I’ll register a backward hook:
mod = nn.Linear(2, 2)

for p in mod2.parameters():
    p.register_hook(lambda x: x/10)

y_hat = mod(x)
y = mod(y_hat)
y.sum().backward()

in that case unwillingly the grads will be twice computed with division by 10, and I wanted only for y gradients will be divided by 10 and going back through y_hat the gradients will remain as is.
So I want to do a manipulation of the same parameters gradients, but they exists twice in the backward graph (if you look with torchviz on the graph of y you will see view functions that create a another second view of the parameter, but that is the same parameter, so if I register a hook, it will be called twice, and I want only once).
That’s why I thought to create a variable as a counter of amount of times I visited that parameter (I know it is 2 at most, so I will increase it every time the hook called, so if it is 1, then I’ll devide the grads by 10 and if it is 2 I will not)



 soulitzer:

We have to register a hook individually for each tensor (aka parameter) in module 2 because if we register the hook directly on the tensor, it triggers when grad is computed for that tensor, and alters only the gradient computation for that individual tensor.


Ok, got you on that. I asked this because I know the chain rule is like Dy/Dz * Dz/Dw …
So I thought although it is individual for each tensor-parameter now we have Dy/(Dz*10)*Dz/(Dw*10) so it will be a mistake to do it for each tensor-parameter because of the math, but you’re saying that not the case and in case of a network of many layers, the correct logic is to register ‘devide_10_hook’ for each layer-parameters, is that correct ?
And the effect in this case will be:
1/10 * (Dy/Dz * Dz/Dw …)