[PyTorch] model.zero_grad() 와 optimizer.zero_grad() 차이

코드를 짜다가 model.zero_grad()와 optimizer.zero_grad()의 차이를 정확히 몰라서 혼용하고 있었습니다. 그래서 헉!하는 마음으로 혹시나 지금까지 했던 학습이 물거품이 될까봐 찾아봤습니다.

 

https://discuss.pytorch.org/t/model-zero-grad-or-optimizer-zero-grad/28426

Model.zero_grad() or optimizer.zero_grad()?

 

일단 PyTorch 공식 포럼 에서는 확인해보면 아래와 같이 사람들이 이야기하더군요.

 

 

두 개의 차이는 거의 없다. 그래서 생각이 다시 들었습니다. 그럼 왜 둘이 구별해뒀지? 조금 더 알아보니 아래와 같은 차이가 있었습니다.

 

optimizer.zero_grad()는 optimizer.__init__(parameters=tensor)에서 optimizer에 register된 parameter 모두에 대해서 tensor.zero_grad()를 호출합니다. model.zero_grad()는 model 아래에 있는 모든 파라미터에 대해 zero_grad()를 수행합니다

 

쉽게 설명하자면 optimizer에 내가 학습하고자 하는 가중치만 zero_grad() 할 경우에는 optimizer.zero_grad() 를 쓰고, 그냥 모델의 모든 가중치를 학습하고자 할 때는 model.zero_grad() 를 써도 된다라고 정리할 수 있겠습니다.