모델의 파라미터로 등록하지 않기 위한 register_buffer

pytorch transformer 구현 중 Positinal Encoding 부분에서 self.register_buffer('pe', pe) 와 같이 사용하는 부분이 있어서 검색

register_buffer 사용

설명

일반적으로 모델 매개 변수로 간주되지 않는 버퍼를 등록하는 데 사용됩니다. 예를 들어, BatchNorm에서 "running_mean"은 매개 변수는 아니지만 상태로써 사용할 수 있다.

Args:
    name (string): name of the buffer. The buffer can be accessed
        from this module using the given name
    tensor (Tensor): buffer to be registered.

사용법

self.register_buffer('running_mean', torch.zeros(num_features))

Reference

https://www.kite.com/python/docs/torch.nn.Module.register_buffer