PyTorch nn.Module · nn.Linear · Linear Layer 직접 구현: 파라미터를 등록하고 순전파를 설계하는 법
`nn.Module`을 상속해 신경망의 구조와 계산 흐름을 정의한다. `nn.Linear`의 내부 계산을 직접 구현하면서 `Parameter`, 초기화, `state_dict`가 어떻게 연결되는지 확인한다. 마지막에는 모든 가중치를 0으로 초기화한 다층 신경망이 왜 학습하지 못하는지 실험한다.
By 송민성
오늘의 목표
nn.Module을 상속해 직접 모델을 만들고, 입력 데이터가 forward를 거쳐 출력으로 변환되는 과정을 설명할 수 있다. nn.Linear 없이도 학습 가능한 선형층을 구현하고 저장된 파라미터를 확인한다.
개념과 직관
1. nn.Module: 모델을 담는 기본 틀
모듈(module, 입력을 받아 계산한 뒤 출력을 만드는 PyTorch 구성 단위)은 PyTorch 모델의 기본 부품이다. 선형층 하나도 모듈이고, 여러 층으로 만든 거대한 신경망도 모듈이다.
nn.Module을 상속할 때 두 메서드를 주로 작성한다.
__init__: 사용할 층과 파라미터를 준비한다.forward: 입력을 어떤 순서로 계산할지 정의한다.
import torch
from torch import nn
class Double(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x * 2
model = Double()
x = torch.tensor([1.0, 2.0, 3.0])
print(model(x))출력:
tensor([2., 4., 6.])model.forward(x)보다 model(x)를 사용한다. nn.Module이 호출 전후에 필요한 PyTorch 기능을 함께 처리하기 때문이다.
그림으로 보면 다음 흐름이다.
입력 x
│
▼
model(x)
│
▼
forward(x)
│
▼
출력2. nn.Linear: 입력 숫자를 섞는 층
선형층(linear layer, 입력에 가중치를 곱하고 편향을 더하는 층)은 다음 계산을 수행한다.
출력 = 입력 × 가중치의 전치 + 편향입력 특성이 3개이고 출력 특성이 2개라면:
입력: [배치 크기, 3]
가중치: [2, 3]
편향: [2]
출력: [배치 크기, 2]예를 들어 한 사람을 [키, 몸무게, 나이]라는 숫자 3개로 표현한다고 하자. 선형층의 출력 하나는 세 숫자에 서로 다른 중요도를 곱한 뒤 모두 더한 값이다.
출력 1 = 키 × w11 + 몸무게 × w12 + 나이 × w13 + b1
출력 2 = 키 × w21 + 몸무게 × w22 + 나이 × w23 + b2가중치(weight, 각 입력을 얼마나 반영할지 결정하는 값)와 편향(bias, 계산 결과를 일정량 이동시키는 값)은 학습을 통해 바뀐다.
layer = nn.Linear(in_features=3, out_features=2)
x = torch.tensor([
[170.0, 65.0, 25.0],
[180.0, 80.0, 30.0],
])
y = layer(x)
print(y.shape) # torch.Size([2, 2])
print(layer.weight.shape) # torch.Size([2, 3])
print(layer.bias.shape) # torch.Size([2])3. Parameter: 학습할 값이라는 표시
파라미터(parameter, 손실을 줄이도록 학습 과정에서 갱신되는 텐서)는 nn.Parameter로 만든다.
value = nn.Parameter(torch.tensor([1.0, 2.0]))
print(value.requires_grad) # Truenn.Module의 속성으로 Parameter를 저장하면 PyTorch가 자동으로 등록한다.
class TinyModel(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.tensor([1.0, 2.0]))
def forward(self, x):
return x * self.weight
model = TinyModel()
for name, parameter in model.named_parameters():
print(name, parameter)일반 텐서를 속성으로 저장하는 것만으로는 학습 파라미터가 되지 않는다.
4. 초기화: 학습 시작점 정하기
초기화(initialization, 학습 전에 가중치와 편향의 시작값을 정하는 과정)는 신경망의 첫 상태를 결정한다.
layer = nn.Linear(3, 2)
nn.init.normal_(layer.weight, mean=0.0, std=0.01)
nn.init.zeros_(layer.bias)가중치를 작은 무작위 값으로 시작하면 각 뉴런이 서로 다른 계산을 한다. 모든 가중치를 같은 값, 특히 0으로 시작하면 여러 뉴런이 똑같이 움직이는 대칭성(symmetry, 여러 뉴런이 같은 출력과 같은 기울기를 갖는 상태)이 생길 수 있다.
5. state_dict: 모델의 학습 상태
상태 사전(state dictionary, 모델이 가진 파라미터와 일부 상태를 이름별로 보관한 사전)은 state_dict()로 얻는다.
layer = nn.Linear(3, 2)
for name, value in layer.state_dict().items():
print(name, value.shape)출력 형태:
weight torch.Size([2, 3])
bias torch.Size([2])state_dict에는 모델 구조를 실행하는 코드가 아니라 학습된 숫자들이 담긴다. 같은 구조의 모델에 이 숫자를 불러오면 같은 계산을 재현할 수 있다.
코드로 직접 해보기
1. nn.Linear의 계산을 먼저 직접 수행하기
import torch
x = torch.tensor([
[1.0, 2.0, 3.0],
[4.0, 5.0, 6.0],
])
weight = torch.tensor([
[0.1, 0.2, 0.3],
[0.4, 0.5, 0.6],
])
bias = torch.tensor([0.5, -0.5])
output = x @ weight.T + bias
print(output)
print(output.shape)@는 행렬 곱셈이다. weight.T는 가중치의 행과 열을 바꾼 전치(transpose)다.
2. Parameter로 선형층 직접 구현하기
import torch
from torch import nn
class MyLinear(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.weight = nn.Parameter(
torch.empty(out_features, in_features)
)
self.bias = nn.Parameter(
torch.empty(out_features)
)
nn.init.normal_(self.weight, mean=0.0, std=0.01)
nn.init.zeros_(self.bias)
def forward(self, x):
return x @ self.weight.T + self.bias
layer = MyLinear(3, 2)
x = torch.tensor([[1.0, 2.0, 3.0]])
print(layer(x))
print(dict(layer.named_parameters()).keys())torch.empty는 값이 정해지지 않은 공간만 만든다. 따라서 생성 직후 nn.init.normal_이나 다른 초기화 함수로 값을 채워야 한다.
3. 직접 만든 층과 nn.Linear 비교하기
두 층의 가중치와 편향이 다르면 출력도 다르다. 공정한 비교를 위해 같은 값을 복사한다.
import torch
from torch import nn
class MyLinear(nn.Module):
def __init__(self, in_features, out_features):
super().__init__()
self.weight = nn.Parameter(
torch.empty(out_features, in_features)
)
self.bias = nn.Parameter(torch.empty(out_features))
nn.init.normal_(self.weight, mean=0.0, std=0.01)
nn.init.zeros_(self.bias)
def forward(self, x):
return x @ self.weight.T + self.bias
torch.manual_seed(0)
custom = MyLinear(3, 2)
builtin = nn.Linear(3, 2)
with torch.no_grad():
builtin.weight.copy_(custom.weight)
builtin.bias.copy_(custom.bias)
x = torch.randn(4, 3)
custom_output = custom(x)
builtin_output = builtin(x)
print(custom_output)
print(builtin_output)
print(torch.allclose(custom_output, builtin_output))마지막 결과가 True라면 두 층이 같은 계산을 수행한다는 뜻이다. torch.no_grad()는 복사 작업을 학습 계산 기록에서 제외한다.
4. nn.Module로 MLP 만들기
다층 퍼셉트론(multilayer perceptron, 선형층과 활성화 함수를 여러 번 연결한 신경망)은 선형층만 연속으로 쌓지 않는다. 중간에 활성화 함수(activation function, 단순한 선형 계산만으로 표현할 수 없는 관계를 학습하게 만드는 함수)를 넣는다.
여기서는 렐루(ReLU, 음수는 0으로 바꾸고 양수는 유지하는 함수)를 사용한다.
ReLU(x) = x가 0보다 크면 x, 아니면 0import torch
from torch import nn
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(2, 8)
self.activation = nn.ReLU()
self.layer2 = nn.Linear(8, 1)
def forward(self, x):
x = self.layer1(x)
x = self.activation(x)
x = self.layer2(x)
return x
model = MLP()
x = torch.tensor([
[0.0, 0.0],
[0.0, 1.0],
[1.0, 0.0],
[1.0, 1.0],
])
print(model(x))
print(model.state_dict().keys())계산 흐름:
입력 2개
│
▼
Linear(2 → 8)
│
▼
ReLU
│
▼
Linear(8 → 1)
│
▼
출력 1개오늘의 실험
모든 가중치를 0으로 초기화한 MLP와 기본 초기화를 사용한 MLP를 비교한다. 데이터는 두 입력이 서로 다를 때만 1인 XOR 문제다.
import torch
from torch import nn
torch.manual_seed(0)
x = torch.tensor([
[0.0, 0.0],
[0.0, 1.0],
[1.0, 0.0],
[1.0, 1.0],
])
target = torch.tensor([
[0.0],
[1.0],
[1.0],
[0.0],
])
class MLP(nn.Module):
def __init__(self, zero_init=False):
super().__init__()
self.layer1 = nn.Linear(2, 8)
self.activation = nn.ReLU()
self.layer2 = nn.Linear(8, 1)
if zero_init:
for parameter in self.parameters():
nn.init.zeros_(parameter)
def forward(self, x):
x = self.layer1(x)
x = self.activation(x)
return self.layer2(x)
def train(zero_init):
model = MLP(zero_init=zero_init)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.BCEWithLogitsLoss()
history = []
for step in range(1001):
prediction = model(x)
loss = loss_fn(prediction, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % 200 == 0:
history.append((step, loss.item()))
with torch.no_grad():
probability = torch.sigmoid(model(x))
result = (probability >= 0.5).float()
return history, probability, result
normal_history, normal_probability, normal_result = train(False)
zero_history, zero_probability, zero_result = train(True)
print("기본 초기화 손실:", normal_history)
print("0 초기화 손실:", zero_history)
print("기본 초기화 확률:\n", normal_probability)
print("기본 초기화 예측:\n", normal_result)
print("0 초기화 확률:\n", zero_probability)
print("0 초기화 예측:\n", zero_result)다음 표를 직접 채운다.
| 초기화 | 0단계 손실 | 200단계 손실 | 1000단계 손실 | 정답 개수 |
|---|---|---|---|---|
| 기본 초기화 | ||||
| 모든 값 0 |
관찰할 내용:
- 기본 초기화에서는 은닉층 뉴런들이 서로 다른 값으로 시작한다.
- 0 초기화에서는 모든 은닉층 뉴런이 같은 출력을 만든다.
- ReLU 입력과 두 번째 층 가중치까지 모두 0이면 첫 번째 층에 유용한 기울기가 전달되지 않는다.
- 균형 잡힌 XOR 데이터에서는 출력 편향만 바꿔도 손실을 충분히 줄일 수 없어 예측이 약 0.5 근처에 머문다.
“0 초기화는 언제나 학습 불가능”이라고 일반화하면 안 된다. 선형층 하나는 0에서 시작해도 입력에 따른 기울기를 받을 수 있다. 문제가 두드러지는 곳은 같은 역할을 하는 여러 은닉 뉴런을 모두 같은 값으로 초기화한 다층 신경망이다.
왜 이걸 배우나 (LLM / Post-training 연결)
LLM도 수많은 nn.Module과 선형층으로 구성되며, 학습된 가중치는 state_dict 같은 형태로 저장하고 불러온다. 사후학습(post-training, 사전학습된 모델을 목적에 맞게 추가 학습하는 과정) 실험에서는 어떤 파라미터가 학습되는지, 초기화와 체크포인트가 올바른지 확인하는 능력이 재현성과 디버깅의 출발점이 된다.
초보자가 막히는 지점
super().__init__()을 빼먹음
부모 클래스인 nn.Module의 초기화가 실행되지 않아 파라미터 등록이 깨질 수 있다.
class Model(nn.Module):
def __init__(self):
super().__init__()forward를 직접 호출함
다음처럼 모델 자체를 호출한다.
output = model(x)행렬 크기가 맞지 않음
nn.Linear(3, 2)는 마지막 차원의 크기가 3인 입력을 기대한다.
layer = nn.Linear(3, 2)
good = torch.randn(4, 3)
bad = torch.randn(4, 5)
print(layer(good).shape) # [4, 2]에러가 나면 입력의 shape, 가중치의 shape, 기대하는 in_features를 먼저 출력한다.
Parameter를 일반 텐서로 만듦
다음 값은 자동으로 학습 파라미터에 등록되지 않는다.
self.weight = torch.randn(2, 3)학습할 값이라면 nn.Parameter를 사용한다.
self.weight = nn.Parameter(torch.randn(2, 3))출력 비교 전에 가중치를 맞추지 않음
MyLinear와 nn.Linear는 처음에 서로 다른 무작위 값을 가진다. 계산식이 같아도 출력은 다르다. 비교 전에 copy_로 같은 가중치와 편향을 넣는다.
알고리즘 (병행 1시간)
- 쉬움: LeetCode 104번
Maximum Depth of Binary Tree
- 깊이 우선 탐색(DFS, 한 경로를 끝까지 방문한 뒤 돌아오는 탐색)으로 왼쪽과 오른쪽 부분 트리의 깊이를 구한다. - 핵심 식: 현재 깊이 = 1 + max(왼쪽 깊이, 오른쪽 깊이)
- 중간: LeetCode 200번
Number of Islands
- 육지를 발견하면 DFS로 연결된 육지를 모두 방문 처리한다. - DFS를 새로 시작한 횟수가 섬의 개수다.
각 문제마다 다음 세 항목을 기록한다.
- 접근: 어떤 상태를 방문 처리했고 DFS가 언제 끝나는가
- 시간복잡도: 입력의 각 노드나 칸을 몇 번 방문하는가
- 틀린 이유: 방문 처리 누락, 경계 검사 오류, 종료 조건 오류 중 무엇이었는가