torch.nn.Module

torch.nn.Module 은 PyTorch의 모든 Neural Network의 Base Class이다.

모듈은 다른 모듈을 포함할 수 있고, 트리 구조로 형성할 수 있다.

 

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

class MyModel(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        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))

 

MyModel을 만들 때, nn.Module을 상속받아서 기본적인 기능들을 사용할 수 있게 만들어준다.

반드시 Add class의 부모 클래스인 nn.Module을 super()을 사용해서 초기화 시켜줘야 한다.

 

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

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

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

 

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

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

>>> @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에 올라가 있습니다.
  
  • cudadevice=Nonedevice=None: 컴퓨터에 CPU는 하나지만 GPU는 여러대 있을 수 있죠? 기본적으로 Torch가 가용할 수 있는 GPU의 첫번째보통0번보통0번가 기본값이지만 따로 설정할 수 있습니다. model.cudadevice=torch.device(′cuda:1′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은 trainmode=Falsemode=False와 똑같습니다.