[PyTorch] torch.no_grad()와 model.eval()의 차이

 PyTorch를 사용하면서 모델의 Evaluation을 진행할 때 Memory leak가 발생하면서 두 함수의 차이점을 분명히 해야한다는 점을 깨달았다. 따라서 본 피드에서는 두 함수의 차이에 대해 공부한 내용을 써내려가보도록 하겠다.

model.eval()
with torch.no_grad():
    for batch in data_loader:
    ...

 

우리는 흔히 모델의 Evaluation 단계를 진행할 때 위와 같은 코드로 흔히 시작한다.

 

이제 두 함수의 차이점을 본격적으로 알아보자.

 

with torch.no_grad()

 이와 같이 torch.no_grad()를 with statement에 포함시키게 되면 Pytorch는 autograd engine을 꺼버린다. 이 말은 더 이상 자동으로 gradient를 트래킹하지 않는다는 말이 된다.

 

 그러면 이런 의문이 들 수 있다. loss.backward()를 통해 backpropagation을 진행하지 않는다면 뭐 gradient를 게산하든지 말든지 큰 상관이 없는 것이 아닌가?

 

 맞는 말이다. torch.no_grad()의 주된 목적은 autograd를 끔으로써 메모리 사용량을 줄이고 연산 속도를 높히기 위함이다. 사실상 어짜피 안쓸 gradient인데 inference시에 굳이 계산할 필요가 없지 않은가? 그래서 일반적으로 inference를 진행할 때는 torch.no_grad() with statement로 감싼다는 사실을 알면 된다.

 

model.eval()

 그럼 여기서 다시 처음 질문으로 돌아와서 위의 torch.no_grad()만 쓰면 되지 않나? gradient 계산 안하고 이제 됐잖아 라고 생각할 수 있다.

 

 맞는 말이지만, model.eval()의 역할은 약간 다르다. 모델링 시 training할때와  inference할때 다르게 동작하는 layer들이 존재한다. 예를 들면, Dropout layer는 학습시에는 동작해야하지만, inference시에는 동작하지 않는 것과 같은 예시를 들 수 있다. BatchNorm같은 경우도 마찬가지다.

 

 사실상 model.eval()는 이런 layer들의 동작을 inference(eval) mode로 바꿔준다는 목적으로 사용된다. (메모리와는 관련이 없다.) 따라서, 우리가 보통 원하는 모델의 동작을 위해서는 위의 두 가지를 모두 사용해야하는 것이 맞다.