torch.nn, torch.nn.Module, torch.nn.Sequential(), torch.nn.ModuleList, torch.nn.ModuleDict

torch.nn

---

These are the basic building blocks for graphs.
(출처 - PyTorch torch.nn 공식 문서)

  torch.nn은 PyTorch 공식 문서에서 적어둔 것과 같이, 그래프를 위한(AI 공부하는 사람들에게는 ML, DL 모델을 만들기 위해 주로 사용하는) 기본적인 블럭들의 모음이다. 레고를 조립하여 작품을 만들 듯, PyTorch에서 블럭들을 만들어두고 이를 torch.nn으로 모아놓았다.

torch.nn 하위 목차

---

  하위 목차로 다양한 레이어와 함수가 속해있다.
Containers, Convolution Layers, Pooling layers, Padding Layers, Non-linear Activations (weighted sum, nonlinearity), Non-linear Activations (other), Normalization Layers, Recurrent Layers, Transformer Layers, Linear Layers, Dropout Layers, Sparse Layers,
Distance Functions, Loss Functions, Vision Layers, Shuffle Layers, DataParallel Layers (multi-GPU, distributed), Utilities, Quantized Functions, Lazy Modules Initialization

  이 중, 모델을 직접 만들기 위해 많이 쓰는 것이 Containers 안의 nn.Module, nn.Sequential, nn.ModuleList, nn.ModuleDict 등이다.

nn.Module

---

• Base class for all neural network modules.
• Your models should also subclass this class.
• Modules can also contain other Modules, allowing to nest them in a tree structure.
  (출처 - PyTorch torch.nn.Module 공식 문서)

  PyTorch의 torch.nn.Module은 공식 문서의 소개와 같이 모든 신경망 모듈을 위한 기반 클래스이다. 대게 많은 모델은 이 클래스를 상속받는다.
nn.Module 클래스는 여러 기능을 한 곳에 모아놓은 바구니이다. 이 바구니에는 다른 바구니 또한 넣을 수 있다. 이때, 기능만 가득 모아둔 바구니라면 basic building block이, basic building block을 모아둔 바구니라면 딥러닝 모델, 딥러닝 모델을 가득 모아둔 바구니라면 더 큰 딥러닝 모델이 된다. PyTorch torch.nn.Module 공식 문서에 따른 예시는 다음과 같다.

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

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

nn.Module 클래스 내부 method

---

  nn.Module에는 많은 내부 method들이 있다. add_module, apply, bfloat16, buffers, children 등을 포함하는데, 종류는 PyTorch 공식 문서를 통해 확인하는 것이 명확하다.     nn.Module PyTorch 공식 문서

nn.Module 클래스 활용

---

더하기 연산 수행 모델 만들기

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

    def forward(self, x1, x2):
        return torch.add(x1, x2)
        
x1 = torch.tensor([1])
x2 = torch.tensor([2])

add = Add()
output = add(x1, x2)
output
>>> tensor([3])

  위와 같이 Add라는 우리가 직접 만드는 클래스는 nn.Module을 상속한다. 이후 __init__에서 초기 설정을 해준다. 추가로 설명하면 forward는 신경망 모델에서 순전파에 해당한다.

nn.Sequential

---

  torch.nn.Sequential()에 내부에 여러 모듈을 원하는대로 배치하여 순차적으로 실행시킬 수 있다. PyTorch nn.Sequential 공식 문서에 나온 예시는 다음과 같다.

model = nn.Sequential(
          nn.Conv2d(1,20,5),
          nn.ReLU(),
          nn.Conv2d(20,64,5),
          nn.ReLU()
        )

  nn.Sequential에는 인자를 위처럼 넣어주어도 되고 이를 OrderedDict 형태로 넣어주어도 된다. 동일한 결과를 얻는다.

model = nn.Sequential(OrderedDict([
          ('conv1', nn.Conv2d(1,20,5)),
          ('relu1', nn.ReLU()),
          ('conv2', nn.Conv2d(20,64,5)),
          ('relu2', nn.ReLU())
        ]))

nn.ModuleList, nn.ModuleDict

---

  torch.nn.ModuleList는 파이썬에서 기초적으로 활용하는 list처럼 모듈을 list에 두고 인덱싱하여 쓸 수 있도록 도와주는 기능이다. 파이썬 list와 다른 점은 리스트를 모은 변수를 불렀을 때 파이썬 list는 값이 사라지는 반면, nn.ModuleList를 통한 Module list는 사라지지 않는다.
  nn.ModuleList는 유용하지만 모듈을 너무 많이 기록해 놓았을 때, 인덱스 번호를 잊어버리면 잘 쓰지 못할 수 있다는 단점이 있다. 이때 torch.nn.ModuleDict를 활용하면 좋다. 모듈의 이름을 key 값으로, 모듈을 value로 설정하여 필요할 때마다 key 값을 통해 불러온다. PyTorch 공식 문서의 예시는 다음과 같다.

//nn.ModuleList
class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linears = nn.ModuleList([nn.Linear(10, 10) for i in range(10)])

    def forward(self, x):
        # ModuleList can act as an iterable, or be indexed using ints
        for i, l in enumerate(self.linears):
            x = self.linears[i // 2](x) + l(x)
        return x

//nn.ModuleDict
class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.choices = nn.ModuleDict({
                'conv': nn.Conv2d(10, 10, 3),
                'pool': nn.MaxPool2d(3)
        })
        self.activations = nn.ModuleDict([
                ['lrelu', nn.LeakyReLU()],
                ['prelu', nn.PReLU()]
        ])

    def forward(self, x, choice, act):
        x = self.choices[choice](x)
        x = self.activations[act](x)
        return x