Module의 상태(state)-Parameter and Buffer

시작하기

이 문서에서는 nn.Module 안에

• 학습 대상인 Parameter와
• 학습 대상은 아니지만 모델 상태로 관리되는 Buffer를 어떻게 추가하는지를 설명함.

이는 Module 객체의
메서드 parameters(), named_buffers(), state_dict() 들이
무엇을 기준으로 동작하는지 훨씬 명확하게 파악할 수 있게 해줌.

---

권장사항

이 글을 보고나서 다음 글의 3. Module의 메서드들 을 다시 한번 읽어볼 것:

2024.04.12 - [Python] - [PyTorch] Custom Model 과 torch.nn.Module의 메서드들.

[PyTorch] Custom Model 과 torch.nn.Module의 메서드들.

 

해당 3장에서 등장하는 다음의 메서드들이 다루는 parameter와 buffer를 설명하는 것이 바로 이 문서의 목적임.

• .parameters()
• .named_parameters()
• .buffers()
• .named_buffers()
• .state_dict()

---

Module에 Parameter와 Buffer 

nn.Module은 forward() 와 함께 모델 (또는 layer, block)이 내부에서 관리해야 하는 상태(state) 도 함께 가짐.

 

이 상태는 크게 다음과 같이 나눌 수 있음.

• Parameter
  • 학습 대상(trainable)인 tensor
  • optimizer가 업데이트할 수 있는 대상
• Buffer
  • 학습 대상은 아니지만, 모델의 상태로 함께 관리해야 하는 tensor
  • state_dict()에 포함될 수 있으며, to(), cuda(), cpu() 호출 시 같이 이동됨

Parameter는 "학습되는 값"이고,
Buffer는 "학습되지는 않지만 모델이 기억해야 하는 값"이라고 보면 됨.

---

Parameter 추가하기

모듈에 학습 가능한 값을 추가하려면 nn.Parameter를 다음과 같이 사용함.

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        # 학습 가능한 파라미터 추가
        self.weight = nn.Parameter(torch.randn(3, 3))
        self.bias   = nn.Parameter(torch.zeros(3))

    def forward(self, x):
        return x @ self.weight + self.bias

위와 같이 nn.Parameter 객체를 self.이름 형태로 대입하면
PyTorch는 이를 해당 모듈의 parameter로 자동 등록함.

 

즉, 다음과 같이 작성하며 parameter가 되고

self.weight = nn.Parameter(...)

 

해당 parameter는 이후 다음과 같은 곳에 자동으로 반영됨.

• model.parameters()
• model.named_parameters()
• optimizer = torch.optim.SGD(model.parameters(), ...)
• model.state_dict()

만약 다음과 같이, 그냥 일반 tensor를 대입하면 parameter로 등록되지 않음!

class BadModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.weight = torch.randn(3, 3)   # parameter로 등록되지 않음

 

이 경우 self.weight는 단지 일반 attribute일 뿐이며,

• model.parameters()에 나타나지 않음
• optimizer가 업데이트하지 않음
• state_dict()에도 기본적으로 포함되지 않음

따라서 학습되어야 하는 값이라면
반드시 nn.Parameter로 감싸야 함.

---

register_parameter() 를 통한 명시적 parameter 등록

Parameter는

• 앞서 다룬 self.weight = nn.Parameter(...) 의 방식으로 추가하는 것이 일반적이나
• register_parameter()를 통해 이름을 문자열로 지정하면서 명시적으로 등록도 가능함

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_parameter(
            "weight",
            nn.Parameter(torch.randn(3, 3))
        )

    def forward(self, x):
        return x @ self.weight

• 이는 내부적으로 parameter registry에 직접 명시적으로 등록하는 방식임.

하지만, 실무에서는 대부분 앞서 다룬 self.weight = nn.Parameter(...) 방식이 더 자주 쓰임.

 

다음을 명심할 것:

• register_parameter()는 가능한 방법 중 하나임
• 하지만, 일반적인 사용자 코드에서는 nn.Parameter(...)를 통한 직접 대입 방식이 더 간단하고 직관적임.

---

Buffer 추가하기

Buffer는 학습 대상은 아니지만 모듈의 상태로 관리해야 하는 tensor를 의미함.

 

대표적인 예는 다음과 같음.

• BatchNorm의 running mean
• BatchNorm의 running variance
• positional encoding
• 추론 시 재사용할 mask
• 클래스 인덱스 매핑용 tensor

 

이러한 값들은 optimizer가 업데이트할 필요(=학습할 필요)는 없으나 다음이 필요함:

• 모델 저장 시 함께 저장되어야 하고
• CPU/GPU 이동 시 함께 이동되어야 하며
• 모듈의 일부로 관리되어야 함

Buffer를 추가하기 위해선
register_buffer()를 사용!!

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("running_mean", torch.zeros(3))
        self.register_buffer("running_var",  torch.ones(3))

    def forward(self, x):
        return x

 

이렇게 등록한 buffer는 다음 특성을 가짐.

• model.buffers(), model.named_buffers()에서 조회 가능
• model.state_dict()에 포함됨
• model.to(device), model.cuda(), model.cpu() 호출 시 함께 이동됨
• optimizer의 학습 대상은 아님

buffer는
"저장되고 이동되지만 학습되지는 않는 tensor" 임.

---

일반 tensor와 buffer의 차이

다음 두 코드는 겉보기엔 비슷하지만 동작이 다름.

 

일반 tensor

class BadModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.mask = torch.ones(3, 3)   # 일반 tensor

 

buffer

class GoodModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("mask", torch.ones(3, 3))   # buffer

 

첫 번째의 self.mask는 일반 attribute일 뿐이므로

• named_buffers()에 나타나지 않음
• state_dict()에 포함되지 않음
• model.cuda() 시 자동 이동되지 않음 (중요)

일반 tensor로 만들어지는 attribute는
그 순간 계산에만 필요한 값,
또는 모듈의 영구 상태로 관리할 필요가 없는 값임.

반면 두 번째의 mask는 buffer로 등록되므로

• 모듈의 상태로 관리됨
• 저장/로딩 대상이 됨
• device 이동 대상이 됨

따라서 "학습은 안 되지만 모델이 반드시 기억해야 하는 tensor"라면
일반 tensor로 두지 말고 register_buffer()로 등록해야 함.

---

buffer에서 persistent=False 옵션

register_buffer()는 기본적으로 등록한 buffer를 state_dict()에 포함시킴.

 

하지만 일시적으로만 필요한 값이라 저장은 할 필요가 없을 경우엔 이럴 때 persistent=False를 사용함.

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("temp_mask", torch.ones(3, 3), persistent=False)

이 경우 temp_mask는

• buffer로 등록되어 device 이동은 함께 되지만
• state_dict()에는 포함되지 않음
• 일반 tensor는 device이동이 안됨을 기억할 것!

즉, 다음의 경우에 유용함:

• 저장은 필요 없음
• 하지만 모듈과 함께 CPU/GPU 이동은 되어야 함

매번 다시 만들기엔 번거롭고,
모델과 같이 움직여야는 하는데,
checkpoint에는 넣고 싶지 않을 때 persistent=False를 사용

 

예를 들어, sinusoidal positional encoding은 공식을 알면 다시 만들 수 있으니 굳이 저장까지는 필요하지 않음.

---

Parameter와 Buffer 비교

항목	Parameter	Buffer	일반 Tensor
의미	모델이 학습하는 값	모델이 기억해야 하지만 학습하지는 않는 값	그냥 계산에 쓰는 일반 값
등록 방법	nn.Parameter(...) 또는 register_parameter()	register_buffer()	그냥 tensor로 생성하여 attribute나 지역변수로 사용
parameters()에 포함	포함됨	포함되지 않음	포함되지 않음
optimizer가 업데이트	가능	불가	불가
state_dict() 포함	포함됨	기본적으로 포함됨	포함되지 않음
to(), cuda(), cpu() 시 이동	이동됨	이동됨	자동 이동되지 않음
named_parameters()에 표시	표시됨	표시되지 않음	표시되지 않음
named_buffers()에 표시	표시되지 않음	표시됨	표시되지 않음
학습 대상 여부	학습 대상임	학습 대상 아님	학습 대상 아님
대표 예	weight, bias	running mean, running variance, 고정 mask	forward() 내부의 중간 계산값, 임시 mask, 입력/출력 tensor

Tensor는 PyTorch에서 가장 기본적인 계산 단위임.

• 이들 중 optimizer가 학습해야 하도록 Module에 등록된 tensor가 Parameter 임
• 학습 대상은 아니지만 모델의 상태로 함께 관리되도록 등록된 tensor가 Buffer임.
• 반면 일반 Tensor는 계산 과정에서 자유롭게 사용하는 값으로, Module이 자동으로 관리하는 대상은 아님.

---

예제코드

import torch
import torch.nn as nn

class DemoModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.weight = nn.Parameter(torch.randn(2, 2))
        self.register_buffer("running_mean", torch.zeros(2))
        self.plain_tensor = torch.ones(2)

    def forward(self, x):
        return x @ self.weight
        
model = DemoModule()

print("named_parameters:")
for name, param in model.named_parameters():
    print(name, param.shape)

print("\nnamed_buffers:")
for name, buf in model.named_buffers():
    print(name, buf.shape)

print("\nstate_dict:")
for key in model.state_dict():
    print(key)

결과는 다음과 같음:

• weight는 named_parameters()에 나타남
• running_mean은 named_buffers()에 나타남
• plain_tensor는 둘 다에 나타나지 않음
• state_dict()에는 weight, running_mean만 들어감

즉, Module이 자동으로 관리해 주는 대상은
그냥 아무 tensor가 아니라, "registered states"들 (Parameters와 Buffer)임을 알 수 있음.

---

요약

nn.Module 내부의 tensor는 모두 동일하게 취급되지 않음.

• 학습되어야 하는 값은 Parameter
• 학습되지는 않지만 모델 상태로 유지되어야 하는 값은 Buffer
• 단순한 계산용 임시 tensor는 일반 attribute

이 구분을 정확히 하고 다음의 method들의 동작을 살펴볼 것:

• .parameters()
• .named_parameters()
• .buffers()
• .named_buffers()
• .state_dict()
• .to()

---

같이보면 좋은 자료들

2024.04.12 - [Python] - [PyTorch] Custom Model 과 torch.nn.Module의 메서드들.

[PyTorch] Custom Model 과 torch.nn.Module의 메서드들.

2024.05.16 - [분류 전체보기] - [DL] PyTorch: state_dict()

[DL] PyTorch: state_dict()