[Pytorch] nn.Module.register_buffer와 nn.Module.register_parameter, nn.Parameters 차이

[Pytorch] nn.Module.register_buffer와 nn.Module.register_parameter, nn.Parameters 차이

 

목차

 

 

torch.nn.Module.register_buffer(name, tensor, persistent)

 

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

Module — PyTorch 1.11.0 documentation

 

 

torch.nn.Module.register_buffer는 parameter가 아니라 말 그대로 buffer를 수행하기 위한 목적으로 활용한다. buffer도 state_dict에 저장되지만 backpropagation을 진행하지 않고 최적화에 사용되지 않는다는 의미이다. 단순한 buffer로써의 역할을 맡는 본 모듈이라고 보면 된다.

 

정리

1) Optimizer가 step 즉, update를 하지 않는다.

2) state_dict에는 저장된다. (나중에 다시 사용하기에 용이)

3) GPU에서 작동한다.

 

참조) 아래 링크의 Buffer Code 예제를 통해 이해를 권장한다.

https://teamdable.github.io/techblog/PyTorch-Module

PyTorch Module

 

torch.nn.Module.register_parameter(name, param)

 

주어진 name을 기반으로 파라미터를 추가하는 함수이다. 이를 통해 parameter를 추가로 설정할 수 있다.

 

1) Optimizer가 update 한다.

2) state_dict에 저장된다.

3) GPU에서 작동한다.

 

nn.Module.register_parameter, nn.Parameters 차이

 

nn.Module.register_parameter takes the tensor or None but first checks if the name is in dictionary of the module. While nn.Parameter doesn’t have such check.

 

예시

import torch
from torch import nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_tensor = torch.randn(1)
        self.register_buffer('my_buffer', torch.randn(1))
        self.my_param = nn.Parameter(torch.randn(1))

    def forward(self, x):
        return x * self.my_param


model = MyModel()

print(model.my_tensor) # tensor([0.2411])
print(model.my_buffer) # tensor([-0.2993])
# Parameter containing: tensor([-0.0892], requires_grad=True)
print(model.my_param) 
# OrderedDict([('my_param', tensor([-0.0892])), ('my_buffer', tensor([-0.2993]))])
print(model.state_dict())

model.cuda()
print(model.my_tensor) # tensor([0.2411])
print(model.my_buffer) # tensor([-0.2993], device='cuda:0')
# Parameter containing: tensor([-0.0892], device='cuda:0', requires_grad=True)
print(model.my_param) 
# OrderedDict([('my_param', tensor([-0.0892], device='cuda:0')), ('my_buffer', tensor([-0.2993], device='cuda:0'))])
print(model.state_dict())
print()

# parameters [('my_param', Parameter containing: tensor([-0.0892], device='cuda:0', requires_grad=True))]
print('parameters', list(model.named_parameters()))

# buffers [('my_buffer', tensor([-0.2993], device='cuda:0'))]
print('buffers', list(model.named_buffers()))
print()

 

참조 자료

 

https://discuss.pytorch.org/t/what-is-the-difference-between-register-buffer-and-register-parameter-of-nn-module/32723

What is the difference between `register_buffer` and `register_parameter` of `nn.Module`

 

https://stackoverflow.com/questions/57540745/what-is-the-difference-between-register-parameter-and-register-buffer-in-pytorch

What is the difference between register_parameter and register_buffer in PyTorch?

 

 

0000