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

 PyTorch로 구현된 코드들에 빠지지 않고 들어가 있는 메서드인 zero_grad(), 공식문서에는 다음과 같이 설명하고 있다.

 이 메서드는 해당 옵티마이저에서 최적화하는 파라미터들의 그라디언트를 모두 0으로 만들어주는 역할을 한다.
 
보통 다음과 같이 쓰인다.

import torch.optim as optim

optimizer = optim.Adam(model.parameters())

...

optimizer.zero_grad()
loss.backward()
optimizer.step()

 
 loss.backward() 를 통해서 loss의 그라디언트를 역전파하는데. 이때 optimizer.zero_grad()를 하지 않으면 기존에 저장되어 있었던 값이 영향을 미치면서 역전파과정이 잘못되어진다. 그래서 반드시 loss.backward() 전에 zero_grad() 메서드를 사용하여 텐서를 0으로 만들어주고 역전파를 시켜야한다.
 
그런데 간혹 다음과 같이 쓰여있기도 하다.
 

import torch.optim as optim

optimizer = optim.Adam(model.parameters())

...

model.zero_grad() ####
loss.backward()
optimizer.step()

 
이렇게 optimizer.zero_grad() 가 아니라 model.zero_grad() 처럼 모델에 대해서 메서드를 적용하는 경우가 있다. (이번에 VAE 구현하면서 막힌 부분을 찾아보다가 당황스러웠다..) 그래서 두 개 모두 적용해본결과 큰 차이를 보이지 않았고 그래서 두 개의 차이가 무엇인지 더 궁금해졌다. 
 
 두 개의 차이는

• 모델 내 파라미터의 그라디언트를 0으로 만드는가
• 옵티마이저에서 최적화하는 파라미터에 대한 그라디언트를 0으로 만드는가

 
 복잡하지 않은 모델의 경우를 살펴보면 대부분 모델 내의 파라미터들은 하나의 옵티마이저로 최적화하게 된다. 이 경우에는 optimizer.zero_grad() 와 model.zero_grad() 가 동일한 역할을 하게된다.
 
 하지만 모델이 복잡해짐에따라 모델 내에 서브 모듈들이 서로 다른 옵티마이저를 사용하는 경우에는 model.zero_grad() 가 훨씬 편한 방법이 될 수 있다. 모든 그라디언트를 0으로 만들기 위해서 옵티마이저 마다 zero_grad() 를 수행하려면 여러 개의 옵티마이저에 대해 zero_grad() 메서드를 모두 호출해야 되기 때문이다.
 
 정리하자면,

• 모델 내의 파라미터와 옵티마이저에서 최적화하는 파라미터가 같다면 optimizer.zero_grad()와 model.zero_grad()가 하는 일은 같다.
• 하지만 모델 내의 파라미터와 옵티마이저에서 최적화하는 파라미터가 서로 다르다면 조금 더 포괄적인 방향인 model.zero_grad() 를 사용하면 된다.