Pytorch 모듈 이해하기

김동환·2023년 3월 14일
0

AI_tech_5기

목록 보기
1/18

Module 이해하기

Parameter

Parameter는 일반적으로 weight나 bias를 설정할 때 tensor 대신에 이용되는 type. parameter을 이용해야만 gradient를 계산하는 grad_fn이 생성되기 때문

self.W = Parameter(torch.ones((out_features, in_features)))
self.b = Parameter(torch.ones(out_features))

Buffer

tensor는 gradient 계산도 x, 값 업데이트도 x, 모델을 저장할 시에 값이 저장되지도 않는다.

parameter는 grad 계산 및 업데이트, 저장 모두 가능

buffer는 모델 저장시 값만 저장되는 type
Batchnorm에서 사용된다고 함

buffer라는 이름으로 등록이 된다.
self.register_buffer('buffer', self.tensor)

값 가져오기

모델 내부에 어떤 module들이 있는지 알고 싶을 때

model.named_childeren

한 단계 하위에 있는 submodule 정보를 준다.

model.named_modules

모든 submodule 정보를 준다.

model.get_submodule(이름)

이름으로 찾는 기능

모델 내부에 Paramter들에 대해서 접근하고 싶을 때. 버퍼도 마찬가지

model.named_parameters()

[ Name ] : ab.b.W1
[ Parameter ]
Parameter containing:
tensor([10.], requires_grad=True)

와 같이 이름과 파라미터에 대한 정보를 얻을 수 있다. 다만 iterator여서 generator를 사용해야

model.get_parameter()

이름으로 파라미터에 접근할 수 있게함

모델 내용 수정하기

docstring

class Model(nn.Module):
    '''doc'''
    def __init__(self):

이런 식으로 doc을 작성하면 model.doc으로 접근을 할 수 있다.

extra_repr

def extra_repr(self):
      return 원하는 내용

repr(model) 출력되는 내용을 수정 가능

Hook 이해하기

pytorch에는 tensor hook & module hook이 존재
register_hook 을 이용해 hook을 등록하면 실행 도중에 hook이 걸린 부분에서 실행이 된다.

class Add(nn.Module):
    def __init__(self):
        super().__init__() 

    def forward(self, x1, x2):
        output = torch.add(x1, x2)

        return output

클래스가 이렇게 있다고 치자

def pre_hook(module, input):
    answer.extend(input)
    pass

def hook(module, input, output):
    answer.extend(output)
    pass
    
add.register_forward_hook(hook)
add.register_forward_pre_hook(pre_hook)

이렇게 되면 forward가 실행 -> pre_hook(x1, x2) 실행 -> output = torch.add() 실행 -> hook(output)이 실행이 된다. 만약 return을 하면 값을 수정하는 것도 가능

def module_hook(module, grad_input, grad_output):
    answer.extend(grad_input)
    answer.extend(grad_output)
    pass
 
 model.register_full_backward_hook(module_hook)

이런 식으로 backward가 실행이 될 때 hook을 실행시켜서 grad값을 체크할 수도 있다.

지금까지는 module에 hook을 걸었는데 tensor에 hook을 적용하려면 register_hook을 쓰면 된다.

hook 활용

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

Apply

module을 입력으로 받아서 해당 모듈의 모든 module을 입력 받아서 처리한다.

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

        self.ab = Layer_AB()
        self.cd = Layer_CD()

    def forward(self, x):
        x = self.ab(x)
        x = self.cd(x)

이 모델에는 2개의 레이어가 있고 각 레이어 안에는 2개의 함수가 들어있다. 이 때 apply를 누르면 A -> B ->AB -> C -> D -> CD -> model 순으로 traversal을 하면서 적용이 된다.

hook, apply 같이 쓰기


def hook(module, input, output):
    module_name = module.__class__.__name__
    output = input[0] @ module.W + module.b
    return output
    
def linear_transformation(module):
    module_name = module.__class__.__name__
    if module_name.split('_')[0] == "Function": module.register_forward_hook(hook)
    
returned_module = model.apply(weight_initialization)
returned_module = model.apply(linear_transformation)

이렇게 하면 Function인 모듈이 실행될 때 forward -> 실행되고 hook이 실행된다. 그렇게 되면 input을 받아서 output을 수정해주는 꼴이 된다. 결국 model -> Function -> forward -> output1(forward에 의한) -> hook -> output2(hook) 이렇게 output이 수정이 된다.

profile
AI Engineer

0개의 댓글