【PyTorch】register_hook的使用

下面是我今天看的代码，即在对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.]])