hook

1. because pytorch The intermediate result of graph calculation will be discarded automatically , So if you want these values, you need to use hook functions . Hook functions include Variable Hook and nn.Module hook , Similar usage .
2. In the use of hook Function time , Its input should not be modified , But it can return a new gradient that replaces the current gradient , namely , Using this function will return a gradient value

register_hook

in the light of Tensor Variable hook function

import torch

grad_list = []
def print_grad(grad):
	grad_list.append(grad)

x = torch.randn(2,1)
y = x+2
y.register_hook(print_grad)
y.backward()

When the whole network carries out reverse transmission , After running to the variable registered by the hook function , The gradient of the variable is saved , call print_grad function , Add gradient to grad_list in

register_forward_hook

For the network layer hook function , Specific visual blog , You can refer to my other document , Visualize specific layers

Hook functions should not modify input and output , And it should be deleted in time after use , To avoid increasing the running load by running the hook every time . Hook function is mainly used to obtain some intermediate results , Such as the output of an intermediate layer or the gradient of a layer .

import torch
model = VGG()
features = torch.Tensor()
def hook(module, input, output):
	#  Copy the output of this layer to features in 
	features.copy_(output.data)

handle = model.layer8.register_forward_hook(hook)
_ = model(input)
#  run out hook Delete after 
handle.remove()

register_backward_hook

First introduced Container The concept of ： When Module Of forward There is only one of the functions Function When , be called Module, If Module Include others Module, be called Container.

stay module Register one on backward hook. The target of this method can only be used in Module On , Cannot be used in Container On .

Every time calculation module Of inputs The gradient of , This hook Will be called

hook(module,grad_input,grad_output)->Tensor or None

Find the gradient of the module , And register_grad similar