cleanUrl: "pytorch-register-buffer"
description: "PyTorch model에서 register_buffer를 사용하는 이유에 대해 알아봅니다."

• PyTorch 모델을 구현할 때 모델 내부적으로 가지고 있으면서, 학습은 필요 없는 tensor들을 사용해야 할 때가 있다.
  • 예를 들면, transformer 구현에서의 causal attention mask 라던지…
• 이 때 흔히 하는 실수가 그 tensor들을 적절한 device (CPU/GPU)로 보내지 않아서 device의 mismatch가 생기는 실수인데, 개인적으로 forward 내부나 모델 구현 내부에서 .to(device)를 call 하는 모양새가 예쁘지 않아서 어떻게 하면 model.cuda() 실행 시에 이러한 non-learnable tensor들을 parameter들과 함께 적절한 device로 옮길 수 있을지 알아보았다.
• nn.Module.register_buffer('attribute_name', tensor)를 이용하면 되더라!
• 아래는 register_buffer 메소드를 실행했을 때의 특징을 정리한 것이다.

<aside> 💡 nn.Module.register_buffer('attribute_name', t)

• 모듈 내에서 tensor t는 self.attribute_name 으로 접근 가능하다.
• Tensor t는 학습되지 않는다. (중요)
• model.cuda() 시에 t도 함께 GPU로 간다. </aside>

예시

model.cuda() 시에 GPU로 이동한다

import torch
import torch.nn as nn

class Model(nn.Module):
	def __init__(self):
		super().__init__()
		
		self.param = nn.Parameter(torch.randn([2, 2]))
		
		buff = torch.randn([2, 2])
		self.register_buffer('buff', buff)

		self.non_buff = torch.randn([2, 2])

	def forward(self, x):
		return x

model = Model()
print(model.param.device)    # cpu
print(model.buff.device)     # cpu
print(model.non_buff.device) # cpu

model.cuda()
print(model.param.device)    # cuda:0
print(model.buff.device)     # cuda:0
print(model.non_buff.device) # cpu

• 위의 예시에서, buffer 에 들어간 tensor는 model.cuda() 시 일반적인 parameter 처럼 GPU로 이동하는 것을 확인할 수 있다.
• 반면 buffer에 넣지 않은 tensor는 그대로 CPU에 남아있다.

state_dict()로 확인이 가능하다

import torch
import torch.nn as nn

class Model(nn.Module):
	def __init__(self):
		super().__init__()
		
		self.param = nn.Parameter(torch.randn([2, 2]))
		
		buff = torch.randn([2, 2])
		self.register_buffer('buff', buff)

		self.non_buff = torch.randn([2, 2])

	def forward(self, x):
		return x

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

OrderedDict([
	('param', tensor([[...]])),
  ('buff', tensor([[...]]))
]

• param과 buff만 state_dict()에서 확인이 가능한 것을 알 수 있다.

Parameter가 아니므로, 당연히 model.parameters()에는 나타나지 않는다.

import torch
import torch.nn as nn
import torch.optim as optim

class Model(nn.Module):
	def __init__(self):
		super().__init__()
		
		self.param = nn.Parameter(torch.randn([2, 2]))
		
		buff = torch.randn([2, 2])
		self.register_buffer('buff', buff)

		self.non_buff = torch.randn([2, 2])

	def forward(self, x):
		return x

model = Model()
for name, param in model.named_parameters():
	print(name, param.data)

param tensor([[...]])

• buff와 non_buff 둘 다 나타나지 않는다.

requires_grad=True로 buffer에 넣어도 parameter로 인식하지 않는다.