[PyTorch] nn.Module에 대한 이해

안녕하세요, 혹시 PyTorch를 쓰시다 보면 습관적으로 nn.Module을 상속 받고 시작하지 않으시나요? 오늘은 타성에 젖은 상속 말고, nn.Module이 뭔지 이해하고 써보면 좋겠습니다. 설명해보겠습니다.

nn.Module 정의

먼저 공식 Docs에 있는 torch.nn.Module부터 읽어보시죠.

PyTorch 공식 문서에서 설명하고 있는 Module

한마디로 모든 모델의 근간이 되는 기본 클래스입니다. 우리가 직접 정의하는 모든 모델은 이 녀석을 상속 받고 시작해야만 forward, backward 등을 편하게 수행할 수 있다는 소리입니다.

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module): # 여기서 nn.Module을 상속 받는다
    def __init__(self):
        super().__init__() # nn.Module과 그 부모가 있다면 전부 상속 받는다
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

반드시 정의가 필요한 함수

이 안에는 너무나 너무나 많은 함수들이 포함되어 있는데요, 정말 중요한 몇가지만 예시로 살펴보겠습니다.

• __init__(): 모델을 만든다고 하면 그 구조를 정의해줘야 하는데 __init__()을 통해 만드는게 규칙입니다.
• forward(): 모델이 어떻게 forward propagation을 할지 정의하는 부분입니다. 똑같이 레이어를 만들더라도, 어떻게 연산을 할 것인지에 따라 완전히 다른 모델이 됩니다.

이 두가지만 있어도 모델은 정의되고 forward/backward를 문제 없이 수행하고 학습할 수 있습니다.

선택적으로 작성하는 함수

다음으로는 많이 쓰는 몇가지 함수들을 예제와 함께 알아보겠습니다.

• apply(fn): 모델의 모든 submodule에 대해 재귀적(recursive)으로 연산을 수행합니다. 이 submodule은 .children() 메소드를 호출했을 때 return으로 얻는 것과 동일합니다.
  • submodule이란? 모델 내 모든 nn.Module을 상속 받는 클래스들을 의미합니다. 이것은 가장 자주 쓰이는 nn.Sequential, nn.ModuleList 등이 전부 포함됩니다. 그 안에는 nn.Linear, nn.Conv2d와 같은 레이어들이 많이 들어가겠죠?
  • apply(fn)는 보통 모델 파라미터의 초기값 설정 이외에는 잘 쓰이지 않습니다...

>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear: # 모델의 모든 submodule에 대해 nn.Linear가 있으면 아래를 수행
>>>         m.weight.fill_(1.0) # fill_(1.0)은 fill의 in-place operation으로, nn.Linear의 weight를 모두 1.0으로 채운다는 뜻
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights) # apply(fn) 적용 방법

• cpu(), cuda(): 모델을 어느 디바이스에 올릴 것인지 결정할 수 있습니다. 기본적으로 모델은 CPU에 올라가 있습니다.
  
  • cuda(device=None): 컴퓨터에 CPU는 하나지만 GPU는 여러대 있을 수 있죠? 기본적으로 Torch가 가용할 수 있는 GPU의 첫번째(보통 0번)가 기본값이지만 따로 설정할 수 있습니다. model.cuda(device=torch.device('cuda:1'))처럼요
• parameters(): 모델의 모든 파라미터를 담은 iterator를 return합니다. 보통 optimizer 선언할 때 이걸 넣게 됩니다.

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

• state_dict(): 모델의 submodule을 dictionary 형태로 반환합니다. 이거 언제 자주 쓰냐면 모델 저장/로드할 때 씁니다.
• train(), eval(): 모델 학습 시작할 때는 model.train(), evaluation 할 때는 model.eval() 쓰시죠? Dropout이나 BatchNorm 같은 레이어들은 학습 시와 평가 시에 다르게 동작하도록 설계되어 있는데요, 모델 내 해당 레이어들을 일일이 모드 변경하지 않고 손쉽게 일괄 변경 시켜줍니다.
  • 참고로 eval()은 train(mode=False)와 똑같습니다.

이 밖에도 register_forward_hook(), register_backward_hook() 등이 있는데, 그것은 나중에 따로 설명하도록 하겠습니다.