cuDNN benchmark 활성화를 통한 최적의 알고리즘 선택

PyTorch 프레임워크 내에는 다양한 CUDA 커널들이 있다. 이로 인해 다양한 환경에서도 최적의 성능을 낼 수 있게 하지만, 어떤 의미로 보면 너무 복잡하게 느껴질 수도 있다.

 

실제로, CNN에서 Conv2D 레이어의 forward pass에 대한 cuDNN 알고리즘만 하더라도 하나가 아니다. 물론, 모두 충분히 빠르다고 할 수 있지만, 그럼에도 현재 실행되는 환경에 더 최적화된 커널이 존재하는 법이다.

파이토치의 백엔드 cudnn API에는 benchmark라는 플래그가 있는데, 이 플래그를 활성화시키면 파이토치에 내장된 auto-tuner가 활성화되면서 가장 최적화된 cudnn 알고리즘을 자동으로 선택해주게 된다.

 

import torch.backends.cudnn as cudnn # type: ignore

cudnn.benchmark = True

위의 코드를 통해서 cudnn algorithm auto-tuner를 활성화시킬 수 있다.

 

추가적인 정보 확인은 https://docs.nvidia.com/deeplearning/cudnn/developer-guide/index.html#tensor-ops-conv-functions-supported-algos에서 가능하다.