nn.Sequential과 모듈 조합 방식

세 가지 컨테이너로 모듈을 조합해

PyTorch에서 모듈을 조합할 때 주로 쓰는 컨테이너는 세 가지야:

nn.Sequential(*modules): 모듈을 순서대로 실행하며 앞 모듈의 출력을 다음 모듈에 넘겨. 별도의 forward()를 작성할 필요가 없어.
nn.ModuleList([modules]): PyTorch가 추적할 수 있는 Python 목록이야. forward()를 직접 작성하고 원하는 방식으로 순회해.
nn.ModuleDict({name: module}): 같은 개념을 사전 형태로 제공해. 다중 헤드 모델이나 분기 구조에 유용해.

'순서대로만 실행할 수 있으면 Sequential, 직접 흐름을 제어해야 하면 ModuleList'라는 기준이 코드의 약 80%에 잘 맞아. Sequential은 읽기 쉽고 print(model)에도 구조가 깔끔하게 보여. ModuleList는 분기, 건너뛰기 연결, 동적으로 정하는 깊이처럼 별도의 제어 흐름이 있을 때 사용해. 예를 들어 설정값 N으로 N계층 Transformer를 만들 때 적합해.

함정: 일반 Python 컬렉션은 쓰면 안 돼

계층을 일반 list나 dict에 저장하면 model.parameters()에 나타나지 않고, .to(device)로 이동하지 않으며, state_dict()에도 저장되지 않아. 더 위험한 건 오류가 나지 않는다는 점이야. 모델은 PyTorch가 볼 수 있는 계층만 학습해. 계층 컬렉션에는 항상 알맞은 nn.Module* 컨테이너를 사용해.

Code

순차: 쉬운 경우·python

import torch.nn as nn

# Plain Sequential
mlp = nn.Sequential(
    nn.Linear(784, 256),
    nn.GELU(),
    nn.Dropout(0.1),
    nn.Linear(256, 10),
)

# Named — useful for inspection and partial freezing
from collections import OrderedDict
mlp_named = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(784, 256)),
    ('act', nn.GELU()),
    ('drop', nn.Dropout(0.1)),
    ('fc2', nn.Linear(256, 10)),
]))

# Now you can refer by name
print(mlp_named.fc1)

ModuleList: 통제 필요할 때·python

import torch
import torch.nn as nn

class FlexibleMLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        # ModuleList — a list PyTorch can see
        self.layers = nn.ModuleList([
            nn.Linear(sizes[i], sizes[i+1])
            for i in range(len(sizes) - 1)
        ])
        self.act = nn.GELU()

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i < len(self.layers) - 1:    # no activation on final layer
                x = self.act(x)
        return x

m = FlexibleMLP([784, 256, 128, 64, 10])
print(sum(p.numel() for p in m.parameters()))  # 234,506

조용한 버그 예: 일반 목록·python

import torch
import torch.nn as nn

class BrokenModel(nn.Module):
    def __init__(self):
        super().__init__()
        # WRONG: regular list. PyTorch can't see these.
        self.layers = [nn.Linear(10, 10) for _ in range(3)]

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

m = BrokenModel()
print(list(m.parameters()))   # [] — empty!
m.to('cpu')                    # silently moves nothing
m(torch.randn(4, 10))          # works at first call (CPU only)

# Fix: use nn.ModuleList instead

세 가지 컨테이너로 모듈을 조합해

함정: 일반 Python 컬렉션은 쓰면 안 돼

Code

External links

Exercise

Progress

댓글 0