데브허브 | DEVHUB | register_buffer() 로 상수 텐서 할당하기 #shortsregister_buffer() 로 상수 텐서 할당하기 #shorts
torch.tensor([값], device='cuda')와 같이 직접 GPU에 텐서를 할당하는 방식은 비동기처럼 보이지만 실제로는 CPU 텐서 생성 후 GPU로 복사되는 동기화 작업이 발생합니다. 🔄
- 이 동기화 과정에서 GPU가 CPU 텐서가 만들어지고 복사될 때까지 대기하게 되어 효율이 떨어집니다. ⏳
- 특히 이 코드가 모델의
forward 함수 내부에 포함될 경우, 매 이터레이션마다 불필요한 복사 작업이 반복되어 훈련 속도가 현저히 느려질 수 있습니다. 🐢
- 이러한 비효율을 피하려면
register_buffer() 함수를 사용하는 것이 좋습니다. ✅
register_buffer() 활용법은 모듈 초기화 시점에 상수 텐서를 한 번 생성하고, register_buffer()를 이용해 인스턴스화 시 GPU로 미리 복사해 두는 것입니다. 🚀
- 이렇게 하면
forward 함수 내에서는 이미 GPU에 복사된 상수 텐서를 계속 재사용하게 되어 불필요한 복사 작업이 발생하지 않아 코드가 효율적으로 작동합니다. ✨