I am trying to find the gradient of Loss function with respect to its neural network (specifically gradient of one node in neural net).
Due to my special type of problem, I only want to compute the gradient with respect to one node in neural network. For example, say we have Loss (L) and neural network parameters theta. Then, I only want to find dL/dTheta_1 where theta_1 is the any ONE node in the neural net.
We typically use grad = torch.autograd.grad(Loss, parameters()) to find gradients of dL/dTheta = [dL/dTheta_1, dL/dTheta_2 ... dL/dTheta_n] but I only want dL/dTheta_1 to reduce the computational cost.
Would it be possible to code it with Pytorch?
Theoretically, I think it is possible to compute only one gradient component but I am not sure Pytorch has the option for that.
Does anyone have a idea on this?
If by a node you mean a single component of a weight tensor, I think it can't be done. Pytorch uses forward graph to perform backward pass, so it can calculate gradients only for graph elements, which theta[0] is not. You could try using autograd to get gradients of outputs of all operations that use theta directly and compute dL/dTheta_1 yourself from them by chain rule. Or you could "chip off" Theta_1 during forward pass from other components and process it separately, which would reduce efficiency of forward pass. But keep in mind that if there are several layers between Theta and L, you can't avoid calculating their full gradients, so it's unlikely you'll save much time.