[Pytorch] hook&apply

In [2]:

from IPython.core.display import display, HTML

display(HTML("<style>.container { width:90% !important; }</style>"))

In [2]:

import torch
from torch import nn
from torch.nn.parameter import Parameter

hook¶

hook 이란 패키지화된 코드에서 다른 프로그래머가 custom 코드를 중간에 실행시킬 수 있도록 만들어 놓은 인터페이스, 코드 중간에 print문으로 내부를 확인하거나 기능을 추가하고 싶을 때 사용한다.

• 프로그램의 실행 로직을 분석
• 프로그램에 추가적인 기능을 제공하고 싶을 때

hook in pytorch¶

Tensor에 적용하는 hook¶

Tensor에는 backward hook만 존재하며, register_hook() 메서드를 통해 hook을 등록하고, _backward_hooks으로 출력한다.

• register_hook - PyTorch 공식 문서

• _backward_hooks

In [6]:

import torch

tensor = torch.rand(1, requires_grad=True)

def tensor_hook(grad):
    pass

tensor.register_hook(tensor_hook)

# tensor는 backward hook만, forward에는 없다.
tensor._backward_hooks

Out[6]:

OrderedDict([(4, <function __main__.tensor_hook(grad)>)])

Module에 적용하는 hook¶

nn.Module에 등록하는 모든 hook은 "__dict__"을 이용해 한번에 확인이 가능하다.

• register_forward_pre_hook - PyTorch 공식 문서
• register_forward_hook - PyTorch 공식 문서
• register_backward_hook - PyTorch 공식 문서
• register_full_backward_hook - PyTorch 공식 문서

In [9]:

class Model(nn.Module):
    def __init__(self):
        super(Model,self).__init__()

def module_hook(grad):
    pass
        
model = Model()
model.register_forward_pre_hook(module_hook)
model.register_forward_hook(module_hook)
model.register_full_backward_hook(module_hook)

model.__dict__

Out[9]:

{'training': True,
 '_parameters': OrderedDict(),
 '_buffers': OrderedDict(),
 '_non_persistent_buffers_set': set(),
 '_backward_hooks': OrderedDict([(7, <function __main__.module_hook(grad)>)]),
 '_is_full_backward_hook': True,
 '_forward_hooks': OrderedDict([(6, <function __main__.module_hook(grad)>)]),
 '_forward_pre_hooks': OrderedDict([(5,
               <function __main__.module_hook(grad)>)]),
 '_state_dict_hooks': OrderedDict(),
 '_load_state_dict_pre_hooks': OrderedDict(),
 '_modules': OrderedDict()}

https://medium.com/the-dl/how-to-use-pytorch-hooks-5041d777f904

backward_hook¶

module 단위의 backward_hook은 module을 기준으로 input, output gradient값만 가져온다.

그래서 Model의 Parameter W의 gradient의 값을 알기 위해선 tensor 단위의 bacwkard_hook을 사용해야 한다.

In [3]:

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.W = Parameter(torch.Tensor([5]))

    def forward(self, x1, x2):
        output = x1 * x2
        output = output * self.W

        return output
    
def tensor_hook(grad):
    answer.extend(grad)
    
answer = []
model = Model()

#tensor 단위 backward_hook
model.W.register_hook(tensor_hook)
x1 = torch.rand(1, requires_grad=True)
x2 = torch.rand(1, requires_grad=True)

output = model(x1, x2)
output.backward() # backward 작동 후에 backward_hook 실행
print(answer)

[tensor(0.0328)]

hook의 활용도¶

• gradient값의 변화를 시각화
• gradient값이 특정 임계값을 넘으면 gradient exploding 경고 알림
• 특정 tensor의 gradient값이 너무 커지거나 작아지는 현상이 관측되면 해당 tensor 한정으로 gradient clipping

apply¶

Python 에 기본적으로 있는 apply처럼(함수를 임의로 만들어 적용) Pytorch에 있는 apply는 module과 submodule과 같이 내부에 있는 모든 모듈들에게 적용되는 custom 함수를 적용하는 것이다.

• pytorch에서 주로 가중치 초기화(Weight Initialization)에 많이 사용된다.

In [1]:

import torch
from torch import nn

@torch.no_grad()
def init_weights(m):
    print(m)
    if type(m) == nn.Linear:
        # 가중치 초기화(Weight Initialization)
        m.weight.fill_(1.0)  
        print(m.weight)

net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
# apply를 통해 적용하는 함수는 module을 입력으로 받는다.
# module.apply(def)
net.apply(init_weights)

Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

Out[1]:

Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)