[PyTorch] Custom Model 과 torch.nn.Module의 메서드들.

Custom Model 만들기

0. nn.Module

 

torch.nn.Module은 PyTorch에서 모든 신경망 모델과 계층의 기반이 되는 클래스임.

• Custom Model (사용자 정의 모델)부터 Built-in Layer(nn.Linear, nn.Conv2d, etc.)까지 전부 nn.Module을 상속받아 만들어짐.
• Model (or Submodule)을 구조화하고 학습 가능한 parameters를 관리하는 데 핵심적인 역할을 함.

---

 

PyTorch에선 Model과 Layer 모두 nn.Module 의 sub-class가 됨: torchvision의 transformer들도 마찬가지임.

 

 

때문에 모든 Custom Model은
nn.Module (또는 그 sub-class)을 상속하여 만들어진다.

 

nn.Module은 다음의 특징을 가짐.

• 생성자(__init__)에서 모델의 nn.Paramter 객체나 nn.Buffer 객체, Sub-Module 객체들을 정의
  • 일반적으로 attribute로 nn.Parameter 객체들이나 nn.Module의 sub-class 객체들을 가짐.
  • nn.Module의 sub-class 객체를 sub-module로 가지는 경우, 반드시 top-level attribute로 추가되어야 함.
• forward 메서드에서 모델의 구성 모듈(=Layer)들을 연결하여 forwared propagation이나 inference등을 tensor 흐름을 지정.
  • __call__를 가진 callable object로서 해당 nn.Module 객체를 직접 호출하여 간접적으로 forward 가 호출되도록 해야함.
  • __call__ 함수에는 hook의 등록 및 autograd 등을 위한 처리들이 forward 호출 이전/이후에 이루어지도록 구현되어 있음.
  • 때문에 forward 메서드를 직접 호출하는 경우 정상적인 동작을 보장하기 어려움. 
• 훈련 중 값이 갱신되는 nn.Parameter 객체들을 관리.
  • parmeters() 메서드와 named_paramerters() 메서드를 이용.
• to(device) 메서드를 통해 CPU에서 동작할지 GPU에서 동작할지를 지정 가능.
• 저장 및 로드를 위한 state_dict() 메서드와 load_state_dict() 메서드를 지원
• 학습 및 추론 모드를 전환
  • DropOut 이나 BatchNormalization 등과 같이 학습 단계와 추론 단계에서 다른 방식으로 실행되는 sub-module이 존재함.
  • 때문에 이를 내부적으로 관리해야 함:train() 메서드로 학습모드로 설정하고, eval() 메서드로 추론모드로 설정.
  • 관련된 attribute가 바로 self.training으로
    PyTorch에서 nn.Module 객체의 현재 상태가 학습 중(training)인지 평가 중(eval) 인지를 나타내는 불리언(Boolean) 값을 가짐.

 

참고로, torchvision에서 transform을 구현할 때의 권장 방식은, nn.Module을 상속받아 forward()에 변환 로직을 정의하고,
랜덤 파라미터는 모듈 내부에서 관리하여 Compose와 자연스럽게 호환되도록 만드는 방식임.

torchvision.transforms 의 모듈의 동작 이해를 위해선 3-4. __call__ 과 __forward__ 부분을 꼭 읽어볼 것.

2025.01.12 - [Python] - [PyTorch] torchvision.transforms 사용법 - transform이란?

[PyTorch] torchvision.transforms 사용법 - transform이란?

---

1. top-level attribute로 sub-module추가.

가장 간단하게 sub-module (layer or activation)을 추가하는 방법임.

 

주의할 것은 top-level attribute가 아닌, list나 dictionary 인스턴스 로 싸고 있는 형태로 추가될 경우, parameters를 통해 제대로 찾지 못하며, 이는 optim 등에서 해당 파라메터나 sub-module의 파라메터를 추적하지 못하는 문제점이나 GPU로 이동시킬 때 해당 파라메터를 제대로 이동시키지 못하는 문제를 일으키게 된다.

 

이에 대한 부분은 다음 URL을 참고할 것.

2024.05.17 - [분류 전체보기] - [DL] PyTorch: nn.ModuleList, nn.ModuleDict, nn.Sequential

[DL] PyTorch: nn.ModuleList, nn.ModuleDict, nn.Sequential

 

다음은 간단한 MLP를 만드는 예제로 각 layer를 Module을 상속한 클래스에서 instance top-level attribute 로 추가하고 있음.

import torch
from torch.nn import Module, init, Linear, Parameter, ReLU
from torch import optim
     
class DsANN (Module): #custom module

    def __init__(self,
                 n_in_f,  # input vector의 차원수.
                 n_out_f, # output vector의 차원수.
                 ):
        super().__init__() # required!

        self.linear0 = Linear(n_in_f, 32)
        self.relu0 = ReLU()

        self.linear1 = Linear(32, 32)
        self.relu1 = ReLU()

        self.linear2 = Linear(32, n_out_f)

        with torch.no_grad():
        #  이 블록 내에서의 연산은 자동 미분(autograd)에서 제외
            # linear0의 bias를 상수 0. 으로 초기화
            init.constant_(self.linear0.bias, 0.)
            
            # Xavier 초기화 방식으로 초기화
            init.xavier_uniform_(self.linear0.weight)

    def forward(self,x):
        x = self.linear0(x)
        x = self.relu0(x)
        x = self.linear1(x)
        x = self.relu1(x)
        y = self.linear2(x)
        return y

 Multi-Layer Perceptron을 간단히 구현함.

• bias와 weights의 초기화 부분은 autograd(자동미분) 그래프에 추가될 필요가 없어서 명시적으로 추가되지 않음을 나타냄.
• 이는 torch.no_grad() 를 이용한 context manager를 통해 이루어짐.
• 생략시 default로 초기화 됨.

2024.11.27 - [Python] - [Py] Context Manager: with statement!

[Py] Context Manager: with statement!

---

2. add_module 로 sub-module 추가.

앞서, instance variable로 처리한 경우와 결과는 같음.,
for문 등의 loop로 구현할 때 사용되는 경우가 많음.

class DsANN (Module): #custom module

    def __init__(self,
                 n_in_f,  # input vector의 차원수.
                 n_out_f, # output vector의 차원수.
                 ):
        super().__init__() # required!

        self.add_module('linear0', Linear(n_in_f, 32))
        self.add_module('relu0', ReLU())
        self.add_module('linear1', Linear(32, 32))
        self.add_module('relu1', ReLU())
        self.add_module('linear2', Linear(32, n_out_f))

        with torch.no_grad():
            # Module의 apply는 특정 함수를 자신의 submodules에 모두 적용 (재귀적).
            self.apply(self.init_weight) 

    def forward(self,x):

        for c in self.children():
            x=c(x)
        return x

        # x = self.linear0(x)
        # x = self.relu0(x)
        # x = self.linear1(x)
        # x = self.relu1(x)
        # x = self.linear2(x)
        # return x

    @classmethod
    def init_weight(cls, module):
        if type(module) == torch.nn.Linear:
            init.kaiming_uniform_(module.weight, mode='fan_in', nonlinearity='relu')
            # init.ones_(module.weight)
            init.constant_(module.bias, 0)

model = DsANN(1,1)
print(model)

---

3. Module의 메서드들.

3-1. Module의 상태 확인하기

.parameters(recurse=True)

optimizer에게 넘겨져 학습을 통해 갱신되어야하는 model의 parameters를 넘겨줄 때 사용됨.

각 parameter의 정보를 볼 때는 .named_parameter()를 이용하면, 이름을 같이 확인할 수 있음.

 

recurse는 재귀적으로 모든 sub-modules의 parameters 까지 반환할지의 옵션으로 기본이 True임.

---

.named_buffers()

buffers는 역시 tensor 객체이나,
학습과정에서 갱신이 필요한 parameters와 달리
학습과정에서 변하지 않는 데이터를 저장하는데 사용된다.

 

.named_buffers()는 모델 내에 정의된 모든 버퍼를 이름과 함께 dictionary 형식으로 반환함.

---

.children()

현재 모델이 가지고 있는 직접적인 sub-moduels (직계 submodules)에 대한 iterator를 반환.

.named_children() 의 경우엔 이름과 함께 dictionary 형식으로 반환한다.

바로 아래의 자식(직계 submodules)에만 접근함.

---

.modules()

현재 모델이 가지고 있는 모든 sub-module에 대해 재귀적으로 반환함.

다음의 예를 확인할 것.

class DoubleLinear(Module):
    def __init__(self, n_in, n_out):
        super().__init__()

        tmp = [(n_in, n_out), (n_out, n_out)]
        for idx, t in enumerate(tmp):
            self.add_module(f'linear{idx}', Linear(*t))
            self.add_module(f'relu{idx}', ReLU())        

    def forward(self,x):
        for c in self.children():
            x=c(x)
        return x

class DsANN (Module): #custom module

    def __init__(self,
                 n_in_f,  # input vector의 차원수.
                 n_out_f, # output vector의 차원수.
                 ):
        super().__init__() # required!

        self.add_module('module1', DoubleLinear(n_in_f,32))
        self.add_module('module2', Linear(32, n_out_f))

        with torch.no_grad():
            self.apply(self.init_weight)

    def forward(self,x):

        for c in self.children():
            x=c(x)
        return x

    @classmethod
    def init_weight(cls, module):
        if type(module) == torch.nn.Linear:
            init.kaiming_uniform_(module.weight, mode='fan_in', nonlinearity='relu')
            # init.ones_(module.weight)
            init.constant_(module.bias, 0)

model = DsANN(1,1)
print(model)

 

여기서 children의 경우는 다음 코드로 확인 가능함.

# for idx, cl in enumerate(model.named_children()):
for idx, cl in enumerate(model.children()):
  print(idx, cl)

 

결과는 다음과 같음.

0 DoubleLinear(
  (linear0): Linear(in_features=1, out_features=32, bias=True)
  (relu0): ReLU()
  (linear1): Linear(in_features=32, out_features=32, bias=True)
  (relu1): ReLU()
)
1 Linear(in_features=32, out_features=1, bias=True)

 

modules() 메서드의 경우는 다음 코드로 확인 가능함.

# for idx, modu in enumerate(model.named_modules()):
for idx, modu in enumerate(model.modules()):
  print (idx, modu)

 

결과는 다음과 같음.

0 DsANN(
  (module1): DoubleLinear(
    (linear0): Linear(in_features=1, out_features=32, bias=True)
    (relu0): ReLU()
    (linear1): Linear(in_features=32, out_features=32, bias=True)
    (relu1): ReLU()
  )
  (module2): Linear(in_features=32, out_features=1, bias=True)
)
1 DoubleLinear(
  (linear0): Linear(in_features=1, out_features=32, bias=True)
  (relu0): ReLU()
  (linear1): Linear(in_features=32, out_features=32, bias=True)
  (relu1): ReLU()
)
2 Linear(in_features=1, out_features=32, bias=True)
3 ReLU()
4 Linear(in_features=32, out_features=32, bias=True)
5 ReLU()
6 Linear(in_features=32, out_features=1, bias=True)

---

3-2. Model의 상태 저장 및 로딩.

.state_dict(desitnation=None, prefix='', keep_vars=False)

현재 모델의 모든 상태 (paramerters and buffers)를 얻어냄.
collections.OrderedDict 로 반환.

• keep_vars 는 기본값이 False로 buffers와 parameters 의 value(값)만을 추출할지를 결정
  • keep_vars=True 인 경우, 값 대신 tensor객체로 데이터 버퍼를 가지고 있는 dictionary가 반환됨.
  • keep_vars=True 인 경우, 메모리 사용량이 커지고, 매우 느리고 복잡한 동작이 이루어지지만. 다음의 장점을 가짐.
    • 모델 디버깅: 모델 상태를 조사하고 특정 매개 변수나 버퍼의 값을 변경해야 하는 경우 유용
    • 모델 커스터마이징: 모델을 불러온 후 특정 매개 변수나 버퍼의 값을 변경해야 하는 경우 유용
    • 모델 저장 및 불러오기 확장: 모델 저장 및 불러오기 프로세스를 확장하고 추가적인 정보를 저장해야 하는 경우 유용
  • 하지만, PyTorch의 버전이 정확히 맞아야만 동작할 수 있는 등의 제한점을 가짐.

2025.04.04 - [Python] - [Py] collections.OrderedDict

[Py] collections.OrderedDict

---

.load_state_dict(state_dict, strict=True)

현재 모델을 "모델상태를 저장한 state_dict"를 이용하여 상태를 설정함.

strict=True 인 경우, 모든 attribute의 이름이 정확히 같아야만 복원됨 (attribute 이름이 key로 사용됨): 기본값은 True임.

---

 

좀 더 자세한 건 다음의 url 참고

2024.05.16 - [분류 전체보기] - [DL] Torch: Save and Load Model

[DL] Torch: Save and Load Model

 

2024.05.16 - [분류 전체보기] - [DL] PyTorch: state_dict()

[DL] PyTorch: state_dict()

---

3-3. Device 지정.

.cuda(device=None)

현재의 모델을 gpu로 복사. (모델의 input tensor들도 gpu로 이동시키는 처리 필요.)

---

.cpu()

현재의 모델을 cpu로 복사. (모델의 input tensor들도 cpu로 이동시키는 처리 필요.)

---

3-4 __call__ 와 forward

 

forward()는 PyTorch의 nn.Module을 상속받은 클래스에서 모델의 실제 연산이 정의되는 부분 임.

 

nn.Module을 상속받아 특정 Class를 개발할 때

대부분 __init_(self, ...)와 forward(self, x) 가 필수적으로 override됨.

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, x):
        print("forward called")
        return self.linear(x)

model = MyModel()

# 일반적인 방식 (추천 방식)
x = torch.randn(1, 10)
output = model(x)  # 내부적으로 model.__call__(x) → model.forward(x) 호출됨

• forward() 는 모델에 입력 데이터를 넣었을 때, 어떤 레이어들이 어떤 순서로 동작하는지를 정의하는 메서드
• forward()는 직접 호출하지 않고, __call__()을 통해 간접 호출하는 것이 권장됨
• nn.Module의 객체는 callable이며, 이 객체를 호출할 때 넘겨지는 모든 arguments는 forward() 메서드에 전달됨.

forward()를 직접 호출하지 않고 __call__()을 통해 호출하는 방식은
PyTorch에서 모델의 추론 또는 학습 시 일반적으로 사용되는 방식임.

 

기본적으로 PyTorch의 nn.Module 클래스는 __call__() 메서드를 오버라이드하여 내부적으로 forward()를 호출하도록 설계되어 있음.

 

 위의 예제처럼 model(x) 처럼 호출하면,

1. nn.Module.__call__() 이 호출됨
2. __call__() 내부에서 다음을 수행:
   • self.training 값에 따라 forward 실행 전후 hook 처리
   • forward() 호출
   • 이 외에도 device 체크, 기타 내부 처리 도 수행됨. 

다음 코드 확인

# torch/nn/modules/module.py, line 483, class: Module
# 간략하게 정리한 버전임.
def __call__(self, *input, **kwargs):
    # --------------------------
    # forward()가 호출되기 바로 전에 실행되는 사용자 정의 코드
	# 예: 입력값을 정규화하거나 모니터링하고 싶을 때 사용
    for hook in self._forward_pre_hooks.values():
        hook(self, input)
 
    # --------------------------
    # 실제 모델의 연산 정의가 수행되는 부분
    # 이 때 학습/평가 모드에 따라 
    # Dropout, BatchNorm 등이 다르게 동작 (self.training이 활용됨)
    result = self.forward(*input, **kwargs)

    # --------------------------
    # forward() 실행이 끝난 후, 출력값에 대한 후처리 가능
    # 예: 특정 레이어의 출력을 저장하거나 시각화할 때
    for hook in self._forward_hooks.values():
        hook_result = hook(self, input, result)
        # ...
    
    # ---------------------------
    # # 일부 문헌에서는 backward_hook도 __call__에 있다고 하지만...
    # # 버전에 따라 다를 수 있으나 backward hook은 
    # # autograd 엔진이 backward를 호출할 때 처리됨.
    # for hook in self._backward_hooks.values():
    #    # ...
 
    return result

 

 

__call__()에서 처리되는 hook 등록이나 JIT tracing, ONNX export 등에서 문제가 될 수 있으므로 forward를 직접 호출해선 안된다.

 

참고로, PyTorch에서 hook란, 모델의 forward 또는 backward 과정 중간에 개입하여 특정 모듈의 입력·출력이나 gradient를 가로채 확인·저장·수정할 수 있게 해주는 메커니즘 임.

자세한 건 다음을 참고: 2025.04.10 - [Python] - [DL] PyTorch-Hook

[DL] PyTorch-Hook

---

같이보면 좋은 자료

https://gist.github.com/dsaint31x/ce49bfbfadc4f95b01684346bfbea76b

dl_modules_methods