Hi, assume we have one input x and we calculate:
x’ = conv_layer(x)
x’’ = conv_layer(x’)
loss2 = criterion(x’’)
and I want to influence the conv_layer parameters only with respect to the gradients graph of x’ (only the path (x—>x’) so the gradient on the path x’ —> x’’ will not be aggregated and influence my optimizer.step action.
The trivial solutions are maybe to create a layer that overrides backward and in the backward path of x’ → x’’ zero the gradients manually instead of aggregating them (or maybe not zero, but substract the added grad value, or just make it the identity backward without any changes to current grad, make it neutral generally) and make it regular on the path I do want influence my optimizer.step.
Another similar solution is maybe to use backward hook to take the gradients from the connection point of the grad paths of x → x’ and x’ → x’’ and then also zero the grads (not sure if I can zero only the part of the unwanted path) and then x’.backward(hooked_grads_from_the_start_of_path_x’->x’’).
Does pytorch supplies some other move convenient and easy way to do so ?
Thank you in advance.