
Tensor indexing 속도 관련 문제 해결 방법

·
아무거나
pytorch tensor의 크기가 클 때 indexing 연산이 생각보다 오래 걸린다. 이를 해결하기 위한 방법을 공유하고자 한다. 코드를 짜면서 크기가 큰 텐서를 인덱싱해야 했다. 인덱싱 한 후 역전파를 하니 학습 시간이 너무 오래 걸려서 어떤 문제인지 찾다가 인덱싱 연산이 생각보다 큰 비용이 발생하는 것을 알게 되었다. CUDA profiling을 통해 각 CUDA kernel이 사용하는 자원을 보았다참고로 CUDA profiling은 아래 처럼 할 수 있다.with torch.autograd.profiler.profile(use_cuda=True) as prof: ''' 프로파일링 하고 싶은 코드 블럭 ''' print(prof.key_averages().table(sort_by..