I am trying to modify the gradients for Multi-task learning. I have a model with multiple modules/layers and I want a list of all gradients in each layer/module. I want to manipulate gradients, let’s say add a similar dimension list to the gradient list computed for simplicity and then want to assign each module/layer the modified gradients. Is there a function to do this, so far I could find I can create a list with the following snippet, how can I assign these gradients back?
def gradient2vec(): gradient_list= for name, param in model.named_parameters(): if param.requires_grad: print name, param.data gradient_list.append(param.grad) return gradient_list def vec2gradient(model, gradient_list): ?????