[Pytorch 기본]nn.Module

nn.Module

nn.Module 클래스는 여러 기능들을 한 곳에 모아놓는 상자 역할을 한다. nn.Module은 빈 상자일 뿐 이를 어떻게 사용할지는 온전히 설계자의 몫이다.

기본적인 클래스에 대한 설명은 파이토치 공식 문서에서 확인할 수 있다.

https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=nn+module#torch.nn.Module 

Module — PyTorch 1.13 documentation

 

기본 구조

x1과 x2를 입력값으로 주면 두 변수의 합을 반환하는 클래스를 만들어 보려 한다.

import torch
from torch import nn

# TODO : Add 모델을 완성하세요!
class Add(nn.Module):
    def __init__(self):
        super().__init__()
        

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


x1 = torch.tensor([1])
x2 = torch.tensor([2])

add = Add()
output = add(x1, x2)

기본적으로 super().__init__을 넣어 상속받는다.

forward 함수에서 x1과 x2를 파라미터로 받아 두 tensor의 합을 return해주었다.

이때 입력값과 출력값은 모두 tensor이다. (중요)

 

torch.nn.Sequential

Sequential은 여러 모듈을 하나로 묶어 실행시키고 싶을 때 사용하는 기능이다.

변수를 파라미터로 받아 인스턴스 변수와 더해주는 Add라는 클래스가 정의되어 있다.

Add를 모듈로 활용하여 여러번 계산하고 싶다.

import torch
from torch import nn

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

    def forward(self, x):
        return x + self.value

class Cal(nn.Module):
  def __init__(self):
    super().__init__()
    self.model = nn.Sequential(Add(3), Add(2), Add(5))
  
  def forward(self, x):
    return self.model(x)

calculator = Cal()

x = torch.tensor([1])
output = calculator(x)

Add를 활용한 Cal이라는 클래스를 만들었다.

nn.Sequential은 여러번의 Add 클래스를 묶어서 순서대로 실행시켜주는 하나의 모델이다.

이때 Add들은 각각 인자로 받은 상수를 self.value에 저장하고 있는 상태로, 앞에서부터 계산된 결과값을 인자로 받아 차례대로 실행된다.

 

nn.ModuleList

파이썬의 list처럼 모듈들을 모아놓기만 하고 그때그때 필요한 것만 골라 사용할 수 있다.

import torch
from torch import nn

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

    def forward(self, x):
        return x + self.value


class Calculator(nn.Module):
    def __init__(self):
        super().__init__()
        self.add_list = nn.ModuleList([Add(2), Add(3), Add(5)])

    def forward(self, x):
      x = self.add_list[0](x)
      x = self.add_list[1](x)
      x = self.add_list[2](x)

      return x
      
x = torch.tensor([1])

calculator = Calculator()
output = calculator(x)

add_list라는 인스턴스 변수에 ModuleList를 저장했다.

forward에서 인덱싱을 통해 필요한 모듈만 골라서 사용할 수 있다.

 

nn.ModuleDict

파이썬의 dict처럼  특정 모듈을 key와 value로 저장할 수 있는 기능이다.

import torch
from torch import nn

# TODO : 다음의 모듈(Module)을 읽고 이해해보세요!
class Add(nn.Module):
    def __init__(self, value):
        super().__init__()
        self.value = value

    def forward(self, x):
        return x + self.value


class Calculator(nn.Module):
    def __init__(self):
        super().__init__()
        self.add_dict = nn.ModuleDict({'add2': Add(2),
                                       'add3': Add(3),
                                       'add5': Add(5)})

    def forward(self, x):
      x = self.add_dict['add2'](x)
      x = self.add_dict['add3'](x)
      x = self.add_dict['add5'](x)
        # TODO : self.add_dict에 담긴 모듈들을 이용하여서
        #        y = ((x + 3) + 2) + 5 의 연산을 구현하세요!
       
      return x


x = torch.tensor([1])

calculator = Calculator()
output = calculator(x)

 

Parameter

Parameter는 클래스 안에서 텐서를 생성해서 파라미터로 저장할 수 있는 기능이다.

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


class Linear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()

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

    def forward(self, x):
        output = torch.addmm(self.b, x, self.W.T)

        return output


x = torch.Tensor([[1, 2],
                  [3, 4]])

linear = Linear(2, 3)
output = linear(x)

output값은 Parameter을 쓰지 않아도 동일하게 계산된다.

 

 

왜 torch.Tensor를 통해서 바로 저장하지 않고 굳이 Parameter를 거칠까?

그 이유는 output tensor를 출력해보면 알 수 있다.

tensor([[4., 4., 4.],
        [8., 8., 8.]], grad_fn=<AddmmBackward0>) # Parameter을 거쳤을 때
tensor([[4., 4., 4.],
        [8., 8., 8.]]) # 바로 Tensor를 생성했을 때

 

이처럼 Parameter를 거쳤을 때만 grad_fn이 생성되어 모델의 학습에 필요한 역전파 과정을 수행할 수 있다.

또한 인스턴스에 state_dict()함수를 통해 확인해보면

linear_parameter.state_dict()

'''OrderedDict([('W', tensor([[1., 1.],
                      [1., 1.],
                      [1., 1.]])), ('b', tensor([1., 1., 1.]))])'''

linear_tensor.state_dict()

 '''OrderedDict()'''

Paramter로 저장했을 경우에만 모델 내에 가중치값이 저장되는 것을 확인할 수 있다.

 

 

Buffer

앞서 tensor은 모델에 값이 저장되지도 않고, 파라미터 업데이트도 진행되지 않는다고 했다.

Buffer은 만약 값이 업데이트 되지 않아도 저장하고 싶은 텐서가 있을 때 활용하는 기능이다.

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


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

        self.parameter = Parameter(torch.Tensor([7]))
        self.register_buffer('buffer', torch.Tensor([7]))
        self.tensor = torch.Tensor([7])


model = Model()

try:
    buffer = model.get_buffer('buffer')
except:
    pass

register_buffer()를 통해 버퍼의 이름과 값을 저장해주고, get_buffer()에서 버퍼의 이름으로 값을 불러온다.

 

세 기능을 정리하자면 다음과 같다.

- "Tensor"
- ❌ gradient 계산
- ❌ 값 업데이트
- ❌ 모델 저장시 값 저장

- "Parameter"
- ✅ gradient 계산
- ✅ 값 업데이트
- ✅ 모델 저장시 값 저장

- "Buffer"
- ❌ gradient 계산
- ❌ 값 업데이트
- ✅ 모델 저장시 값 저장