신경망 코드를 작성할 때 가장 흔히 발생하는 실수들 (feat. 카파시)

들어가면서

최근 카파시의 영상과 트윗(X)을 보면서 크게 공감하는 부분이 많아 이 글을 작성하게 되었다. 인공지능연구원에서 처음 들어갔을 때 가장 힘들었던 것이 이론과 코드를 매칭하는 작업이었다. 역전파의 개념을 이해하고 있고 계산할 수는 있는데, 이를 코드로 작성하는 것은 또 다른 문제였다. 딥러닝의 구조와 개념을 알지만 이걸 코드로 작성하는 것도 또 다른 문제였다. 특히 아카데미에서는 그리 큰 모델을 다루지 않아서 주피터 노트북(?) 정도면 충분했지만, 연구원에서는 주피터 노트북은 무슨.. 수많은 py 파일과 불안한 눈빛, 그리고 이를 정신없이 따라가는 나.. 당시에는 어떤 파일이 핵심 코드이고 어떤 파일이 보조 역할을 하는지조차 제대로 파악하지 못했다..
특히 인공지능 코드는 일반적인 개발 코드와는 다르게 오류가 발생하지 않아도 오류가 날 수 있다(!). 오히려 명확한 에러 문구가 등장하는 것이 럭키할 지경이다. AI 코드는 단순한 디버깅 이슈보다는 (물론 에러도 무진장 많이 난다) 수식이나 로스 함수, 아키텍처에서 논리적인 오류가 발생하고 이를 잡아내지 못하면 모델은 점점 멍청해진다.
이번 글에서는 신경망 코드를 작성할 때 가장 흔히 저지를 수 있는 실수에 대해 다뤄보려고 한다.

카파시 X 게시물 링크

위는 카파시가 언급한 4가지 실수이다. 이와 함께 내가 겪은 다른 실수도 덧붙여 설명할 계획이다.

1. you didn’t try to overfit a single batch first

이는 sanity check 과정과 같다. 즉, 모델이 정상적으로 학습하는지 검증하기 위해 작은 데이터셋에서 로스가 0으로 빠르게 수렴하는지 확인하는 과정이다. 작은 데이터셋에서조차 overfitting이 되지 않는다면, 큰 데이터에서도 제대로 학습되지 않을 가능성이 매우 높다. 따라서 모든 데이터를 학습시키기 전에, 우선 단일 배치 혹은 매우 작은 데이터셋에서 모델이 정상적으로 학습하는지 빠르게 확인해야 한다. 이 과정에서 loss가 줄어들지 않는다면 아키텍처, learning rate, optimizier 등에 문제가 있을 확률이 높다. 이런 sanity check를 거쳐 모델이 정상적으로 overfit 할 수 있다는 것을 확인한 후, 본격적으로 전체 데이터셋을 학습시키는 것이 바람직하다.

2. you forgot to toggle train/eval mode for the net

파이토치 모델은 train과 eval 모드를 자동으로 전환하지 않기 때문에, 학습과 평가 단계에서 각각 train()과 eval()을 직접 설정해주어야 한다. train 모드에서는 dropout이 활성화되고, 일부 뉴런이 랜덤하게 비활성화된다. 또한 현재 배치의 평균/표준편차를 사용해 정규화한다. eval 모드에서는 dropout이 비활성화되어 모든 뉴런이 사용되며, 학습된 전체 데이터셋의 평균/표준편차를 사용한다. 즉, train과 eval에서 사용하는 평균과 분산이 다르기 때문에 batch normalization을 사용하는 모델에서 eval 모드를 설정하지 않으면 배치 크기에 따라 성능이 이상해질 수 있다.

3. you forgot to .zero_grad() before .backward()

역전파를 실행하기 이전에 반드시 .zero_grad()를 통해 파라미터의 기존 gradient를 초기화 해주어야한다. 왜일까? 간단하게 더하기 연산의 backward 코드를 구현해보자.

def __add__(self, other):
    out = Tensor(self.data + other.data)

    def _backward():
        self.grad = out.grad
        other.grad = out.grad
    out._backward = _backward

위 코드에서 self 변수가 한 스텝에서 사용되어 backprop을 했다고 치자. 그럼 self 변수에 gradient가 구해진다. 그런데 이 변수가 하나의 연산에서만 사용되는 것이 아니라, 다른 연산에서도 또 사용된다면 어떻게 될까? 이전 gradient 값이 사라지고, 새로운 gradient가 계속해서 overwrite될 것이다. 즉, 이전 연산들의 gradient 정보가 전부 손실된다. 이를 방지하기 위해 파이토치에서는 gradient를 계산할 때 +=를 통해 누적하는 방식을 사용한다.

def __add__(self, other):
    out = Tensor(self.data + other.data)

    def _backward():
        self.grad += out.grad  # 기존 gradient를 유지한 채 누적
        other.grad += out.grad
    out._backward = _backward

이렇게 하면 같은 변수가 여러번 사용되더라도 각 연산에서 발생한 gradient의 영향을 모두 반영할 수 있다.그리고 이것이 바로 backward 전 .zero_grad()를 해주어야 하는 이유이다. 한 스텝에서 gradient를 업데이트하고 난 후, 다음 스텝에서 backward를 다시 수행하면 어떻게 될까? 동일 변수가 여러번 사용될 것을 대비해 gradient를 +=로 누적해왔으므로, 다음 스텝에서 아무런 조치 없이 다시 backward를 수행하게 되면 이전 스텝의 gradient가 계속 쌓여서 엉뚱한 값이 된다. 따라서 새로운 gradient를 계산하기 전에 반드시 .zero_grad()를 호출하여 이전 스텝에 대한 gradient를 초기화한 후 backward()를 수행해야 한다.

4. you passed softmaxed outputs to a loss that expects raw logits

파이토치에서는 여러 손실함수를 제공하지만, 손실 함수에 따라 입력으로 기대하는 값이 다르다. 어떤 함수는 raw logits을 입력으로 기대하고, 어떤 함수는 확률값(softmax나 sigmoid가 적용된 값)을 입력으로 기대한다. 예를 들어, torch.nn.CrossEntropyLoss() 는 내부적으로 softmax 연산을 적용하기 때문에, 직접 softmax를 적용한 후 이 함수를 사용하면 잘못된 결과가 나올 수 있다. 반면 torch.nn.NLLLoss()는 softmax 적용 후 log값을 입력으로 받는다. 따라서 F.log_softmax()로 변환한 값을 전달해야 한다. 나는 이진분류에서 이런 실수를 한 적이 있었다. raw logit을 받는 F.binary_cross_entropy_with_logits()가 아닌 F.binary_cross_entropy() 에 raw logit을 넘긴 적이 있다. 소프트맥스 대신 시그모이드를 활용할 뿐 다른 것은 위 내용과 같다. 즉, 손실함수를 사용할 때 입력값이 raw logits인지, 확률값인지 반드시 확인해야 한다.

5. pretrained 모델을 사용할 때 주의점

이건 개인적으로 겪었던 경험담인데, 당시 BERT 모델을 활용해 태스크를 수행중이었다. 그런데 아무리 학습을 시켜도 계속해서 같은 임베딩 값이 나오는 것이다. 이전에는 이런 일이 없었기에 그간 수정했던 config들과 모든 사항들을 검토했었다. 처음에는 발전시켰던 아키텍처 구조 혹은 데이터가 문제일 것이라고 생각해서 그 쪽만 계속해서 팠었는데, 아무리 봐도 명확한 오류를 찾을 수 없었다. 몇몇 데이터가 라벨링이 잘못 되어있긴 했지만 퍼센트로 따지면 극히 적었기 때문에 모델에게 이 정도의 혼란을 주지는 못했을 것이라 판단했다. 결론적으로 문제의 원인은 freeze_pretrained = False 옵션이었다… 당시 BERT가 우리 모델에 좀 더 적합하게 튜닝될 수 있지 않을까 하는 생각에 freeze를 하지 않고 학습을 진행했는데, 혹시나 하는 마음에 True로 변경했더니 정상적으로 동작했다. 지금 생각해보면 가능성있는 원인은 두 가지다. 1) BERT는 이미 매우 잘 훈련된 모델이기에 우리 모델의 학습률이 너무 커서 가중치가 깨졌을 가능성 2) LayerNorm이 학습 도중 특정 조건에서 업데이트 되지 않아 gradient 흐름에 문제가 발생했을 가능성 혹은 이 밖의 다른 문제일 수도 있다… pretrained 모델을 활용할 때는 freeze 여부와 학습률을 신중하게 조정해야 한다. 특히, 처음에는 가중치를 freeze한 상태에서 학습을 시작하고 이후 필요에 따라 서서히 풀어가는 것이 안전하다.

마치며

처음 신경망 모델을 작성할 때는 아키텍처를 잘 설계하고 파이토치로 뚝딱뚝딱 쌓기만 하면(물론 이 과정도 결코 쉽지는 않았다) 로스가 잘 줄어드는 모델이 탄생하고 그게 다 인줄 알았다. 처음에는 이만큼만 해도 성장했다는 느낌이 들었는데, 점점 모델이 고도화 될 수록 프레임워크에 의존하면 안되겠다는 생각을 많이 했다. 모델을 개발할 때는 한 단계씩 검증 과정을 거치면서 작성하는 것이 좋다. 일단 이것저것 갖다 붙여놓고 로스가 줄어든다고 와~ 하면 안 된다는 것… 나중에 모든 퍼즐이 맞춰졌을 때 제대로 동작하는 모델을 만들려면, 과정 하나하나를 꼼꼼히 따져야 한다. 특히 pytorch에서 제공하는 클래스들을 가져다 쓰기만 하고 그 안의 매커니즘을 제대로 모르면 발전할 수 없다고 느꼈다. 예를 들어, 알고보니 모델에 문제가 있었고 backpropagation 수식을 직접 수정해야한다면? 구체적인 매커니즘을 모른 채 프레임워크의 힘을 빌려 쓰기만 한다면 모델을 정교하게 개선하는 것은 불가능에 가깝다. 또한 신경망 모델의 경우 잘못되고 있는 줄도 모르고 혼자 조용히 (…) 망가지고 있는 경우가 많기 때문에 좀 더 꼼꼼히 들여다볼 필요가 있다. 위에 서술한 실수들도 글로만 읽으면 ‘당연한거 아냐?’ 싶을 수 있고, 나 또한 그랬다. 하지만 실무에서 직접 코드를 작성하다보면 큰 흐름에 정신이 팔려 이런 사소한 부분들을 놓치게 되기도 하더라.. (그렇게 모델은 또 조용히 망가진다..) 나중에 알고보면 저런 사소한 실수들이 결국 나비효과를 일으켜 모델의 결과가 이상하게 나왔던 적도 적지 않았다.
처음에는 단순히 파이토치로 모델을 완성하는 것만으로도 만족스러웠지만, 모델이 고도화될수록 이론과 구현을 더 깊이 이해해야만 제대로 된 개선이 가능하다는 걸 깨닫게 되었다. 또 이건 다른 말이지만 Typing 등의 파이썬에 대해서도 더 깊이 알아야겠다는 생각이 들었다. 프레임워크가 편리함을 제공하는 도구일 뿐, 문제를 해결해주는 마법이 아니라는 사실을 잊지 말아야겠다고 다짐했다. 결국, 내가 만드는 모델을 내가 직접 이해할 수 있어야 한다. 그리고 가장 중요한 건, 사소한 것이라도 놓치지 않도록 끊임없이 점검하고 개선하는 자세가 아닐까 싶다. 이 글이 나처럼 모델이 조용히(?) 망가진 경험이 있는 사람들에게 조금이나마 도움이 되었으면 좋겠다.