[Python] How i just caculate a part of grad using make_functional in pytorch

func_model, func_params = make_functional(self.model)
def fm(x, func_params):
fx = func_model(func_params, x)
return fx.squeeze(0).squeeze(0)
def floss(func_params,input):
fx = fm(input, func_params)
return fx

per_sample_grads =vmap(jacrev(floss), (None, 0))(func_params, input)


In this code, per_sample_grads includes all the grad of a network model by input some data. However, What i want to get is a small part of grad. And get the all the gradient is some kind of waste in my codes. I want to save some time in this auto_grad step. How can i do to acheive that.

How can I compute the gradients for only part of the model to save time?"

Continue reading...