self.register_buffer(), 학습하지 않을 파라미터라면? (tensor와 명백하게 다른 점 2)

🤔 Problem

오랜만에 PyTorch 관련 글입니다. 최근에는 Generative Model 쪽을 공부하면서 DDPM을 구현하다가 PyTorch의 새로운 기능을 발견했는데요. 바로 오늘 글의 주제가 되는 self.register_buffer()입니다. 본 포스트는 예전에 작성한 포스트들 중에 'nn.Parameter(), 이걸 써야 하는 이유가 뭘까? (tensor와 명백하게 다른 점)'라는 포스트의 후속 편이 되기도 합니다. 

 

이전 포스트의 내용을 간략하게 리뷰해 보면 '모델 내에서 단순히 torch.tensor()를 통해 선언한 텐서는 학습의 대상이 되지 못하고, 이를 명확하게 모델 내 학습을 하는 파라미터로 정의하기 위해서는 nn.Parameter()로 감싸서 추가적으로 선언해야 한다.'라는 내용이었습니다.

 

위 내용에 덧붙여서 'torch.tensor()는 학습을 하지는 않지만, nn.Parameter()처럼 모델 내에 속해 있다고 볼 수는 없다.'라는 내용을 먼저 전달드리고 싶습니다. '모델 코드 안에 선언을 해두었는데 모델에 속해있지 않다는 게 무슨 말이지?'라고 보실 수 있습니다만, 제가 전달하고자 하는 '모델 내에 속해있다'라는 말의 정의는 model.state_dict()와 큰 연관이 있습니다.

 

model.state_dict()는 모델의 weight를 저장할 때 사용하는 메서드입니다. 아래와 같이 torch.tensor()를 통해 정의한 weight가 포함된 모델을 저장하려고 메서드를 사용했을 때는 해당 weight가 보이지 않습니다. 즉, 모델의 구성요소로 인정하고 있지 않습니다.

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.tensor = torch.randn(3, 3)
        
        self.param = nn.Parameter(torch.randn(3, 3))
    ...

model = Model()
print(model.state_dict())

[Output]
OrderedDict({'param': tensor([[ 0.3108,  0.5101, -0.8290],
        [ 0.3511, -0.4658,  0.2131],
        [ 0.6602, -1.0786, -0.3299]])})

 

하지만, 프로젝트를 진행하거나 논문을 보다 보면, 학습을 하지 않음에도 불구하고 모델 내 구성 요소로 취급되는 텐서, 행렬, 스칼라를 보게 됩니다.(ex: DDPM의 alpha, beta)

 

그전에 계속 읽기 불편하지 않게 '학습을 하지 않는 모델의 구성 요소'라는 것을 나타낼 단어가 필요합니다. 학계에서는 이를 단순히 Non-trainable Parameter, Non-learnable Parameter, Fixed Parameter라는 단어로 정의를 합니다. 본 포스트에서는 편하게 Fixed Parameter라는 단어를 사용하겠습니다.

 

다시 돌아와서 이 Fixed Parameter는 단순히 torch.tensor()로 정의되지는 못 한다는 사실을 알게 되었습니다. 그렇다면, PyTorch에서는 이에 대해 어떠한 방법을 제시하고 있을까요?

---

😀 Solution

PyTorch에서는 Fixed Parameter를 Buffer라고 정의합니다. 최근 논문들에서도 Buffer라는 단어를 일컫는 경우들이 종종 있기도 합니다. 그리고, 이 Buffer를 사용하기 위해서 PyTorch는 self.register_buffer()라는 기능을 제공하고 있습니다. 사용하는 방법은 정말 간단합니다. 

 

아래와 같이 torch.tensor()로 선언한 다음에 self.register_buffer()에 할당하여 선언할 수 있습니다. 첫 번째 인자에는 Buffer가 사용할 이름이 되고, 두 번째 인자에는 Buffer로 사용할 텐서가 들어가게 됩니다. 그리고, forward() 함수에서 다른 Parameter처럼 똑같은 접근이 가능한 걸 볼 수 있습니다.

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
        buff = torch.randn(3, 3)
        self.register_buffer('buff', buff)
    
    def forward(self, x):
        print(self.buff)
        return x

 

self.register_buffer()의 장점은 이것뿐만이 아닙니다. 디바이스 간 이동의 자율성이 높다는 장점이 하나 더 있습니다. torch.tensor()로 선언한 경우에는 model = Model.to(device)를 하더라도 다른 디바이스로 넘어가지 않는 불편함이 있었습니다만, self.register_buffer()에서는 이를 해결해주고 있습니다.

 

즉, Buffer로 정의하는 방법은 2가지 장점이 있다고 볼 수 있습니다.

1. 학습을 하지 않지만, 모델의 구성요소로 인정할 수 있다. (가중치로 저장할 수 있다.)
2. 모델의 디바이스 이동 작업이 생길 때 번거로운 작업이 불필요해졌다.

그리고, model.parameters()나 model.name_parameters()로 파라미터를 확인할 수 있는 것처럼 model.buffers()나 model.named_buffers()와 같은 메서드로 Buffer를 확인할 수 있습니다.

 

또한, Buffer라는 특성을 통해 Distribution(torch.distributions)도 관리가 가능한지 실험해 보았으나, Buffer는 Tensor type만 담을 수 있는 것으로 확인하였습니다.

 

마지막으로, 2가지 장점을 확인할 수 있는 코드와 Output을 같이 올리면서 글을 마무리하겠습니다.

# Expr 1
import torch
from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.tensor = torch.randn(3, 3)
        
        self.param = nn.Parameter(torch.randn(3, 3))
        
        buff = torch.randn(3, 3)
        self.register_buffer('buff', buff)
    
    def forward(self, x):
        print(self.buff)
        return x

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Model().to(device)

print("===========================================")
print(f"tensor: \n{model.tensor}")
print(f"{model.tensor.device}")

print("===========================================")
for name, parameter in model.named_parameters():
    print(f"{name}: \n{parameter}")
    print(f"{parameter.device}")

print("===========================================")
for name, buff in model.named_buffers():
    print(f"{name}: \n{buff}")
    print(f"{buff.device}")

print("===========================================")
print(f"model state_dict(): param O, buff O, tensor X")
print(model.state_dict())

# Expr 2
# Distribution은 buffer에 register 안 된다.
from torch.distributions.normal import Normal

class Model2(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.gaussian = Normal(loc=torch.zeros(3),
                               scale=torch.ones(3))
        
        gaussian2 = Normal(loc=torch.zeros(3),
                           scale=torch.ones(3))
        
        self.register_buffer('gaussian2', gaussian2)
    
    def forward(self, x):
        sample = self.gaussian.sample()
        
        sample2 = self.gaussian2.sample()
        
        print(f"sample 1 device: {sample.device}")
        print(f"sample 2 device: {sample2.device}")
        
        return x

model2 = Model2()
model2(torch.randn(3, 3))

'''
(.venv) E:\DDPM>python register_buffer_test.py
===========================================
tensor: 
tensor([[ 0.5325,  1.3698, -1.2790],
        [-0.5546,  0.3236,  0.6196],
        [ 2.1521,  1.8287,  1.0600]])
cpu
===========================================
param:
Parameter containing:
tensor([[-1.5823, -1.0639, -0.9007],
        [-0.8665, -0.0151,  0.8802],
        [ 0.3128,  2.1903, -0.2867]], requires_grad=True)
cpu
===========================================
buff:
tensor([[-0.0590,  0.4823,  2.1716],
        [-0.6110,  0.1420,  1.5730],
        [ 1.2040, -0.0654, -0.4525]])
cpu
===========================================
model state_dict(): param O, buff O, tensor X
OrderedDict({'param': tensor([[-1.5823, -1.0639, -0.9007],
        [-0.8665, -0.0151,  0.8802],
        [ 0.3128,  2.1903, -0.2867]]), 'buff': tensor([[-0.0590,  0.4823,  2.1716],
        [-0.6110,  0.1420,  1.5730],
        [ 1.2040, -0.0654, -0.4525]])})
Traceback (most recent call last):
  File "E:\DDPM\register_buffer_test.py", line 64, in <module>
    model2 = Model2()
             ^^^^^^^^
  File "E:\DDPM\register_buffer_test.py", line 52, in __init__
    self.register_buffer('gaussian2', gaussian2)
  File "E:\DDPM\.venv\Lib\site-packages\torch\nn\modules\module.py", line 566, in register_buffer
    raise TypeError(
TypeError: cannot assign 'torch.distributions.normal.Normal' object to buffer 'gaussian2' (torch Tensor or None required)
'''

---

📂 Reference

https://pytorch.org/docs/stable/generated/torch.nn.Module.html

Module — PyTorch 2.5 documentation

https://www.ai-bio.info/pytorch-register-buffer

PyTorch에서 register_buffer를 써야하는 이유