[Pytorch] 3. torch 라이브러리 내부 구조 분석 (Convolution)

[주의] 개인 공부를 위해 쓴 글이기 때문에 주관적인 내용은 물론, 쓰여진 정보가 틀린 것일 수도 있습니다!

피드백 부탁드립니다. (- -)(_ _) 꾸벅

---

1, 2편에서 대충 어떻게 initialize되고 C extension이 되는지에 대해 동작 원리를 파악했다. 나는 현재 하려는 건 Convolution layer의 동작을 파악하여 channel-wise든 spatial-wise든 좀 더 분산 처리하는 방향의 모델을 연구하고자 하는 것이었기 때문에 Conv layer의 동작 방식을 이해하는게 필요하였다. 그래서 이번에는 Conv layer 위주의 분석을 진행하였다.

 

torch/nn/modules/conv.py

 

 실제 모델을 만들 때 쓰이는 Convolution layer는 /torch/nn/modules/conv.py 에 정의가 되어 있다.

conv.py에는 ConvNd의 자식 클래스로 Conv2d가 정의되어 있다. Conv2d에서는 컨볼루션 레이어의 파라미터들을 초기화해준다. 모델에서 conv layer가 호출될 때 __call__ 함수가 실행되며, 거기에선 forward 함수가 실행이 되기 때문에 forward를 따라 bottom-up approach를 하였다.

 

conv.py의 Conv2d 클래스

 

Conv2d forward 함수에서는 다시 F.conv2d를 반환해준다. F는 functional을 import한 것이므로 functional 모듈을 찾아보았다.

 

functional.py의 conv2 정의 부부

 functional 모듈에서는 conv2d에 관한 정의가 _add_docstr() 밖에 없었다. _add_docstr은 함수를 정의한 것이 아닌, 미리 정의된 함수에다가 doc string을 추가하여 문서화한 것이기 때문에 source code를 bottom-up 방식으로 찾을 수 없었다. 나는 anaconda버전으로 배포한 pytorch를 설치한 것이기 때문에, source code가 폴더에 없었다. github에서 pytorch 소스코드를 다운받아 찾아보기로 한다. (깃헙 : github.com/pytorch/pytorch)

 

 약간의 검색을 통해 pytorch/aten/src/Aten/native/Convolution.cpp 가 컨볼루션 레이어의 소스코드임을 알아내었다.

 

pytorch/aten/src/Aten/native/Convolution.cpp

conv2d의 함수는 다시 at::convolution의 함수를 return한다.

 

pytorch/aten/src/Aten/native/Convolution.cpp

at::convolution 함수는 다시 at::_convolution 함수를 return하는데, 중간에 벤치 마크 모드 활성 여부인 benchmarkCuDNN, Randomness제어인 Deterministic 등을 bool 타입으로 제어해주는데, 나는 그것보단 Conv 함수의 소스코드를 찾는게 우선이기 때문에 다루지 않았따.

 

pytorch/aten/src/Aten/native/Convolution.cpp

컨볼루션 레이어의 꽤 밑에 단의 소스코드인 듯하다. ConvParams의 구조체를 통해 각 파라미터를 포현해주고 있다.

그 밑에는 각 상황별로 호출하는 output 매서드가 다른 것 같다. 길어서 정리해보자면

 

_convolution()

테스트해 본 결과 CPU를 이용했을 때 use_mkldnn이 Trun였고, GPU를 이용했을 때 use_cudnn이 True 였다.

정리하자면

 

다음과 같이 불려지는 것 같다.

 

더 깊은데 까지의 수정은 필요없을 수도 있으니 일단은 여기까지만 분석하였다.