언제 custom function이 필요한가

새 연산을 실험하거나, 기본 autograd가 원하는 방식으로 동작하지 않을 때 torch.autograd.Function을 쓸 수 있다. 이 계층은 Python 수준에서 forward와 backward 계약을 직접 써보는 좋은 출발점이다.

장점

  • 빠르게 아이디어를 실험할 수 있다
  • backward semantics를 명시적으로 정의할 수 있다
  • 이후 C++/CUDA extension으로 내려가기 전 인터페이스를 검증할 수 있다

주의할 점

  • forward에서 무엇을 저장할지 신중해야 한다
  • backward의 shape, dtype, device semantics를 맞춰야 한다
  • 성능은 Python 경계 때문에 제한될 수 있다

즉 custom function은 최종 고성능 구현이라기보다 연산 계약을 먼저 분명히 하는 계층으로 보는 편이 좋다.

다음 글에서는 tensor lifetime과 allocator를 본다. PyTorch 메모리 문제를 다루려면 autograd만큼 allocator와 buffer 재사용 감각도 중요하다.