nn.Module: 모든 모델의 공통 기반

메서드 두 개로 모든 걸 만들어

PyTorch의 모든 신경망 구성 요소는 선형 계층 하나부터 매개변수 700억 개짜리 Transformer까지 nn.Module을 상속해. 하위 클래스를 만드는 방식은 이 프레임워크에서 가장 중요한 Python 관례이며, 계약은 두 가지로 간단해:

__init__을 재정의해. 항상 super().__init__()을 먼저 호출하고, self.layer1 = nn.Linear(10, 20)처럼 자식 모듈과 매개변수를 self에 할당해. PyTorch는 __setattr__을 가로채 이 객체들을 자동으로 등록해.
forward(self, x, ...)를 재정의해 계산을 정의해. 다만 직접 호출하지는 마. model.forward(x)가 아니라 모듈 인스턴스를 함수처럼 model(x)로 호출해야 해. 그래야 __call__을 거치며 등록된 훅을 실행하고, 자식 모듈의 학습·평가 모드를 적용하고 autograd 처리 흐름을 올바르게 이어 줘.

자동으로 얻는 기능

model.parameters(): 중첩 깊이와 관계없이 학습 가능한 모든 텐서를 순회해.
model.named_parameters(): 같은 텐서에 점으로 구분한 경로 이름을 붙여 순회해.
model.to(device): 모든 매개변수와 버퍼를 지정한 장치로 옮겨.
model.train() / model.eval(): 드롭아웃과 BatchNorm처럼 모드에 따라 달라지는 자식 모듈의 동작을 전환해.
model.state_dict() / load_state_dict(): 직렬화할 수 있는 모든 매개변수와 버퍼를 사전 형태로 저장하고 복원해.

이 기능들은 모두 nn.Module의 등록 체계에 의존해. 자식 모듈은 self에 할당하고, 컬렉션에는 일반 Python 목록 대신 nn.ModuleList를 쓰고, 직접 만든 학습 가능한 텐서는 nn.Parameter로 감싸야 해. 이 규칙을 빼먹으면 자동 기능이 오류 없이 조용히 멈춰.

Code

최소 nn.Module: 기준 하위 클래스·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self, in_dim=784, hidden=128, out_dim=10):
        super().__init__()                    # ALWAYS first
        self.fc1 = nn.Linear(in_dim, hidden)
        self.act = nn.ReLU()
        self.fc2 = nn.Linear(hidden, out_dim)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

model = TinyMLP()
print(model)
# TinyMLP(
#   (fc1): Linear(in_features=784, out_features=128, bias=True)
#   (act): ReLU()
#   (fc2): Linear(in_features=128, out_features=10, bias=True)
# )

무료 도구: 매개변수, 장치, 모드·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP()

# Parameter iteration
for name, p in model.named_parameters():
    print(name, p.shape)
# fc.weight torch.Size([4, 10])
# fc.bias   torch.Size([4])

# Total parameter count
total = sum(p.numel() for p in model.parameters())
print(f"Params: {total:,}")  # Params: 44

# Move once, everything follows
model = model.to('cpu')        # or 'cuda' / 'mps'
print(next(model.parameters()).device)

# Mode switching
model.train()                  # default mode
model.eval()                   # affects Dropout / BatchNorm

왜 model(x), 절대 model.forward(x) 아님·python

import torch
import torch.nn as nn

class HookedLinear(nn.Linear):
    pass

m = HookedLinear(4, 2)

# Register a hook that prints the OUTPUT shape after each forward
def hook(module, inputs, output):
    print(f"Hook fired: out shape = {output.shape}")

m.register_forward_hook(hook)

x = torch.randn(3, 4)

m(x)               # Hook fired: out shape = torch.Size([3, 2])
# m.forward(x)     # NO hook fires! Bypasses __call__.

# This is why everyone calls model(x). Bypassing __call__ skips:
#   - registered forward / backward hooks
#   - __torch_function__ dispatch
#   - some compile / quantization machinery

Exercise

784차원 입력 → 256차원 은닉층 → 128차원 은닉층 → 10차원 출력으로 이어지고 각 계층 사이에 ReLU가 들어가는 TinyMLP를 만들어 봐. named_parameters()와 전체 매개변수 수를 출력한 뒤, 784*256 + 256 + 256*128 + 128 + 128*10 + 10을 직접 계산한 값과 맞는지 확인해.

메서드 두 개로 모든 걸 만들어

자동으로 얻는 기능

Code

External links

Exercise

Progress

댓글 0