下面是我今天看的代码,即在对x求导时,将梯度保存下来。
def save_gradient(self, grad):
self.gradients.append(grad)
def __call__(self, x):
outputs = []
self.gradients = []
#for name, module in self.shared_model._modules.items():
#print(module)
#x = module(x)
for name, module in self.model._modules.items():
x = module(x)
#print(name)
if name in self.target_layers:
#print(name)
x.register_hook(self.save_gradient)
outputs += [x]
register_hook的作用:即对x求导时,对x的导数进行操作,并且register_hook的参数只能以函数的形式传过去。
下面举个栗子:
#正常求导情况
v = torch.randn((1, 3), dtype=torch.float32, requires_grad=True)
z = v.sum()
z.backward()
print(v.grad)
tensor([[1., 1., 1.]])
对导数进行2倍操作:
v = torch.randn((1, 3), dtype=torch.float32, requires_grad=True)
z = v.sum()
v.register_hook(lambda grad: grad*2)
z.backward()
print(v.grad)
tensor([[2., 2., 2.]])