Pytorch:hooking,nn.Module,autograd에 대한 고찰

부스트캠프 ai테크 2기 첫 P스테이지인 이미지 분류 대회가 열렸습니다.  이번 대회의 주제는 'COVID 19 Mask 탐지' 였습니다. 과거, 이미지 분류 대회(2021 Programmers 머신러닝 Dev-Matching)를 참가했을때에는 , pretrained_model을 keras에서 불러와서 , 실험-피드백-실험 과정을 반복하면서 성적을 내는데에 집중을 하였습니다. 하지만, 이번 대회는 기간이 몇시간이 아닌 2주였고, pretraiend_model을 이용하는 것은 많이 해보았기 때문에 , 성적이 저조하더라도 직접 모델을 제작하여 , 대회에 참여하고 싶었습니다. 1주차에서는 주어진 미션을 수행하였고, 2주차에는 1주차에서 나온 결과물에서 발생한 문제점을 해결하고자 시도했습니다. 

1주차 모델의 학습 결과를  보았을 때 , train_loss 함수에 변화가 없음을 알 수 있었습니다. 따라서, 저는 

weight initialization을 안해서 vanishing gradient문제라고 처음에는 결론을 내렸습니다. activation에 대해 gradient계산을 할때 0이 되거나 0에 가까운 값이 전달 되는것도 우려되었습니다. 그러다가, 혹시 모델작성에서 실수를 하여서, 제가 의도한대로 parameter들이 학습이 되지 않은 부분이 있나라는 가능성도 떠올랐습니다. 2가지 문제를 체크하기 위해서는 backward를 진행하면서 계산하는 gradient 값을 전부 찍어보는것으로 해결이 가능했기에 모든 gradient값을 찍는 방법을 찾아보았습니다.

우선적으로 고려한것은 autograd.backward의 작동방식을 알아보는 것입니다. autogard.backward를 여태까지는 남들이 썻으니까 썻습니다. 저는 이것에 대해서 정체도 모르는 것에 저를 맡기는 것은 엔지니어로서 툴에 휘둘리는 것이기 때문에  backward 의 작동방식을 이해해서 사용 주체가 저 자신이 되어서 제가 의도하는 대로 작동을 하는지 확인을 해보았습니다. Elliot Waite 이라는 분의 Youtube 동영상 중 pytorch autograd 동영상을 보고 공부했습니다. Google Software Enginner이셧던 분인데 워너비 회사의 Google Engineer 답게 안의 내부 호출대상 까지 세세하게 설명이 되어있었습니다. backward를 호출하게 되면 , requiers_grad=True 인 variable들에 대해서 backward Graph 를 만들고 , 이 노드들에 대해서 , chain_rule을 이용해서 gradient 값을 output으로 내놓거나 input으로 받습니다.  backward 노드들이 계산을 한 후 호출한 variable들에게 gardient 값을 전달하는 데 이 때 특정 variable값을 gradient 계산 후에 inplace-operation으로 변경해주면 , 노드들이 gradient 값을 다시 계산해서 전달해주는것을 처음알았습니다. 이 inplace-operation 을 사용할 때는 주의를 기해야합니다. 그리고 , 이 값들은 변경됫다는 표시를 version이라는  attribute로 표기하였습니다.  

backward를 호출할 때 backward 노드들이 어떻게 만들어지는지 알았다면, 이 노드들에서 계산되어지는 gradient 값을 출력해서 눈으로 확인을 해야합니다. 모든 파라미터들에 대하여 , grad 라는 attribute를 통해서 grad를 출력 할 수 있습니다. 2개~3개 라면 직접 print를 parameter마다 grad를 출력해도 되지만, layer가 100개 넘는 모델에 대하여 이를 일일이 찍을 수는 없습니다. 따라서, 저는 이 작업을 게으르게 하기 위해서 hooking 이라는 기능을 사용하기로 했습니다. Pytorch에서는 forward, backward hooking을 제공해주고 있습니다. hooking이란 특정 기능을 호출하고 나서나, 호출하기 전에  미리 정해둔 동작을 실행함을 의미합니다. forward hooking을 이용하면 중간의 feature map을 찍을 수 있기 때문에 , feature map을 확인할 일이 생기면 forward hooking을 이용하면 됩니다. backward hooking은 backward를 호출하고나서 호출되는 함수입니다. 이 함수는 gradient 값을 계산후 수정 할 수 있는 capsulization을 위반할 수도 있는 신경을 써서 써야하는 기능이였습니다. 이게 신기했던게 , 각 backward 노드마다 +,* 같은 연산에 대해 variable이 2개 밖에 연결되지 않았습니다. 덧셈으로 연결된 경우 backward hooking으로 직접 weight값을 바꾸었을 때 , 덧셈으로 연결된 모든 variable의 gardient에 영향을 주었습니다. backward hooking 은 모든 NN.moudle에 대해서 호출할 때submodule에 대하여 grad_in,grad_out,module 을 input으로 받는 custom 함수를 정의하고 backward hook이 custom 함수를 매번 호출하는 방식입니다. 저는 이 3개를 모두 출력하도록 함수를 작성하였고, 모든 weight를 찍었지만, weight의 개수가 출력하고 살펴보기엔 너무 많은양이었습니다. 그래서, custom 함수 내부에서 새로운 함수를 호출하게 하였습니다. nan값인지 , None값인지 체크해서 해당되면 nan,NOne을 출력하도록 했고 , 그렇지 않다면, 평균값을 출력하도록 했습니다.  이 평균값을 통해 Variable(parameter)들이  chain_rule을 통해 weight를 input ,output으로 주고받는지 확인 했습니다. 확인결과, backward Node에서의 input , output 관계에는 문제가 없었습니다. 하지만, 제가 forward안에서 정의한 Sequential Module의 gradient가 계산되지 않고 skip되었습니다. 저는 이게 왜 이렇게 되었는지 고민하였습니다. 이는, forward라는 함수 스택영역에서 Sequential Module을 정의하였기에 호출이 끝나면서 , 메모리에 올라와있던 Sequential Module이 사라진 것입니다. 메모리에 올라와 있지않다면, gradient를 계산하기 위해서 접근할수 없기 때문에 , gradient를 계산하려면 무조건 메모리에 올려야하므로, class that inherits from nn.Module 의 멤버변수로 하위 모듈을 작성해야 함을 배웠습니다.  

계산값을 출력하고, 모듈정보를 출력하는 과정을 거쳤지만, 개괄적으로 빠르게 파악하기 위해선 시각화도 필수라고 생각을 했습니다. 저는 처음에 backward의 retain_graph 인자가 backward graph를 저장해서 tensorboard로 logging 할 수 있게 하는 인자인줄 알았습니다. 하지만, 이는 backward가 끝나고 사라지는 backward graph를 저장하게 해주는 옵션이였습니다. 시각적으로 활용하는 정보라고 착각한 저는, backward의 개념에 대한 부족으로 인해 여기에 대해서 삽질을 하다가, torchiviz라는 Tool을 알게 되었습니다.  이 Tool은 오픈소스로 , Pytorch의 backward Node graph를 시각화 해주는 도구였습니다. 파이토치 공식 문서에는 아직 이러한 기능을 지원해 주지 않으므로 대부분 이 툴을 사용한다고 하였습니다.  이것으로 layer가 100개넘는 모델을 시각화하여 원하는 부분을 좀더 빠르게 파악할 수 있도록 도움을 주었습니다.