Parameter, Buffer, State Dict

Module 안 두 종류 tensor

Parameter (nn.Parameter) — learnable. model.parameters() 에 등장. optimizer 가 update.
Buffer (self.register_buffer(name, tensor) 로 등록) — non-learnable 이지만 stateful. 예: BatchNorm 의 running mean/var, position encoding, attention mask, teacher model 의 EMA weight.

둘 다 state_dict() 에 저장되고 둘 다 .to(device) 따라 이동. 차이는 optimizer 가 손대는지만.

state_dict — universal serialization 형식

model.state_dict() 가 parameter / buffer 이름을 그 tensor 에 매핑하는 OrderedDict 반환. load_state_dict(d) 가 복원. 이게 model 저장의 추천 형식 — framework-version-tolerant: model 클래스가 진화해도 이름과 shape 가 일치하는 한 (또는 mismatch 허용 위해 strict=False) 새 architecture 에 weight 로드 가능.

Freezing

parameter freeze (optimizer 가 update 안 하게) 위해선 p.requires_grad = False. transfer learning 시작 — pretrained backbone freeze, 새 head 만 train. optimizer 는 여전히 parameter 봐 (model.parameters() 에 있음) 만 gradient 없으니 update 할 게 없음.

Code

parameter 와 buffer 세기·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 — save 와 load·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]))

transfer learning 위한 freeze·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']):,}")

Parameter, Buffer, State Dict

Module 안 두 종류 tensor

state_dict — universal serialization 형식

Freezing

Code

External links

Exercise

Progress

댓글 0