본문 바로가기

연구/pytorch

[pytorch] register_buffer 란?

반응형

register_buffer 로 layer를 등록하면 어떤 특징이 있는가?

1. optimizer가 업데이트하지 않는다.

2. 그러나 값은 존재한다(하나의 layer로써 작용한다고 보면 된다.)

3. state_dict()로 확인이 가능하다.

4. GPU연산이 가능하다.

 

따라서 네트워크를 구성함에 있어서 네트워크를 end2end로 학습시키고 싶은데 중간에 업데이트를 하지않는 일반 layer를 넣고 싶을 때 사용할 수 있다.

 

참고 사이트 : 

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

반응형