[PyTorch] autograd hook이란? Tensor.register_hook(hook)

torch.Tensor.register_hook — PyTorch 1.13 documentation

Tensor.register_hook(hook)
hook(grad) -> Tensor or None

 

autograd는 텐서를 자동으로 미분하는 Python 라이브러리이다. DNN에서 노드(node)는 텐서를 나타내고 엣지(edge)는 해당 텐서에서 작동하는 함수를 나타내는 계산 그래프로 정의할 수 있다. autograd를 사용하여 모델의 순방향 패스를 지정하면 backpropagation 알고리즘을 사용하여 loss 함수와 관련하여 모델 매개변수의 그래디언트를 자동으로 계산한다.

autograd hook은 텐서에 define하고 attach할 수 있는 함수로, 해당 텐서의 그래디언트가 계산될 때마다 autograd에 의해 호출된다. 이는 특정 텐서의 변화도를 기록, 수정 또는 다른 텐서를 업데이트하는 데 사용하는 등의 작업을 수행하려 할 때 유용할 수 있다. 모델의 동작을 사용자가 직접 지정할 수 있도록 하는 강력한 함수가 될 수 있지만 코드의 전체 구조를 복잡하게 만들 수 있으므로 주의해서 사용해야 한다.

autograd hook을 정의하려면 텐서의 register_hook 메서드를 사용할 수 있다. 아래는 예시이다.

import torch

x = torch.tensor(1.0, requires_grad=True)

def my_hook(grad):
    print('Gradient:', grad)

x.register_hook(my_hook)

y = x ** 2 + x + 1
y.backward()
y = x * 20
y.backward()

출력 결과는 아래와 같다.

각 y를 x에 대해 미분하면 각각 2x + 1, 20이 나오며 이에 x=1을 대입하면 3, 20이 나오기 때문에 위와 같은 결과가 나온다.