매개변수, 버퍼, 상태 사전

모듈 안에는 두 종류의 텐서가 있어

매개변수(nn.Parameter): 학습 과정에서 옵티마이저가 갱신하는 텐서야. model.parameters()에 포함돼.
버퍼(self.register_buffer(name, tensor)로 등록): 학습 대상은 아니지만 모델 상태로 보존하는 텐서야. BatchNorm의 누적 평균과 분산, 위치 인코딩, 어텐션 마스크, 교사 모델의 EMA 가중치가 대표적이야.

둘 다 state_dict()에 저장되고 .to(device)를 따라 장치도 옮겨. 차이는 옵티마이저가 갱신하느냐뿐이야.

state_dict: 널리 쓰이는 직렬화 형식

model.state_dict()는 매개변수와 버퍼의 이름을 해당 텐서에 매핑한 OrderedDict를 반환해. load_state_dict(d)로 복원할 수 있어. 모델을 저장할 때 이 형식을 권장하는 이유는 프레임워크와 모델 클래스가 변해도 비교적 견고하기 때문이야. 이름과 모양만 일치하면 새 구조에 가중치를 불러올 수 있고, 일부 불일치를 허용하려면 strict=False를 사용할 수 있어.

매개변수 동결

옵티마이저가 매개변수를 갱신하지 않게 하려면 p.requires_grad = False로 설정해. 전이 학습에서는 보통 사전 학습된 백본을 동결하고 새 헤드만 학습해. 동결한 매개변수도 model.parameters()에는 남아 있어 옵티마이저가 볼 수 있지만, 기울기가 없으므로 갱신되지는 않아.

Code

매개변수와 버퍼 세기·python

import torch
import torch.nn as nn

class DemoModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 4)
        self.bn = nn.BatchNorm1d(4)
        # A custom non-learnable buffer
        self.register_buffer('class_centroids', torch.zeros(4, 4))

    def forward(self, x):
        return self.bn(self.fc(x))

m = DemoModel()

# Parameters — learnable
print('PARAMETERS')
for n, p in m.named_parameters():
    print(f"  {n}: shape={tuple(p.shape)}, requires_grad={p.requires_grad}")

# Buffers — non-learnable state
print('BUFFERS')
for n, b in m.named_buffers():
    print(f"  {n}: shape={tuple(b.shape)}")

state_dict: 저장과 불러오기·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()
torch.save(model.state_dict(), '/tmp/model_weights.pt')

# Load
model2 = TinyMLP()
sd = torch.load('/tmp/model_weights.pt', weights_only=True)
model2.load_state_dict(sd)

# Confirm
for k in sd:
    print(k, torch.equal(sd[k], model2.state_dict()[k]))

전이 학습을 위한 동결·python

import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze the whole backbone
for p in model.parameters():
    p.requires_grad = False

# Replace the head — new params default to requires_grad=True
model.fc = nn.Linear(model.fc.in_features, 5)

# Optimizer should only see the trainable params
import torch.optim as optim
optimizer = optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=1e-3,
)
print(f"Trainable params under optimizer: {sum(p.numel() for g in optimizer.param_groups for p in g['params']):,}")

모듈 안에는 두 종류의 텐서가 있어

state_dict: 널리 쓰이는 직렬화 형식

매개변수 동결

Code

External links

Exercise

Progress

댓글 0