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

참고 https://yuevelyne.tistory.com/10

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

Pytorch 로 모델 학습을 진행한 뒤 모델의 Evaluation 을 진행 할 때 

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

위와 같이 model.eval() 을 사용을 하고 torch.no_grad() 를 뒤이어 또 사용해주는 경우를 보신적이 있으실 겁니다. 

 

위 그림은 pytorch 공식 문서에서 model.eval() 함수의 설명을 발췌한 것 입니다.

 

eval() 함수는 해당 모델의 모든 레이어가 evaluation mode 에 들어가게 해줍니다. 즉, Dropout, BatchNorm 등의 학습시에만 필요한 기능을 비활성화 시켜주는 역할을 합니다. 

 

torch.no_grad() 는 역전파 시 필요한 gradient 계산 context를 비활성화 해주는 역할을 합니다. 예측(추론;Inference) 혹은 Validation(검증)을 할 때는 gradient 계산을 하지 않습니다. 그래서 이 함수를 사용해 줌으로써 PyTorch 의 autograd engine( gradient 를 계산해주는 context ) 를 비활성화 시켜서 더 이상 gradient 를 트래킹(추적)하지 않게 됩니다. 따라서 필요한 메모리가 줄어들고 연산속도가 증가하게 됩니다.