[Pytorch] optimizer.zero_grad() vs nn.Module.zero_grad() 차이

[Pytorch] optimizer.zero_grad() vs nn.Module.zero_grad() 차이

 

일반적인 NN에서는 다음과 같은 순서로 훈련이 진행됩니다.

Mymodel = model()
loss_fn = torch.nn.~Loss()
optimizer = torch.optim.SGD(Mymodel.parameters())


for input, target in dataset:
  optimizer.zero_grad()
  output = model(input)
  loss = loss_fn(output, target)
  loss.backward()
  optimizer.step()

일반적으로 zero_grad는 훈련 단계에서 역전파 이전에 gradient를 0으로 만들 때 사용합니다. 그런데 해당 함수는 모델에서 또한 사용이 가능함을 확인할 수 있습니다. (nn.module 역시도 zero_grad를 사용함을 확인할 수 있습니다.)

Mymodel = model()
loss_fn = torch.nn.~Loss()
optimizer = torch.optim.SGD(Mymodel.parameters())

for input, target in dataset:
  Mymodel.zero_grad()
  output = model(input)
  loss = loss_fn(output, target)
  loss.backward()
  optimizer.step()

그렇다면 두 개의 차이는 무엇일까요? 

 

optimizer.zero_grad()와 nn.Module.zero_grad() 모두 변화도를 0으로 만드는 것이 같습니다. 정확히는 optim.zero_grad()는 변화도 0을 위해 nn.Module.zero_grad()의 모든 parameter들을 호출합니다. 이에 대한 식을 확인하기 위해서는 우리가 설정한 optimizer 코드를 봐야합니다.

 

현재 위에 설정한 optimizer 코드내용은 torch.optim.SGD(Mymodel.parameters())입니다. 즉, optimizer는 Mymodel의 모든 parameter를 가지고 backpropagation된 loss를 다시 재정립하는 과정을 거치는 것입니다. 때문에 optimizer.zero_grad()와 nn.Module.zero_grad()는 위 상황에서는 같은 것으로 보아도 무방합니다.

 

그렇다면 구분해서 사용할 때 유용하다고 볼 수 있는데 그 때는 언제일까요?

 

optimizer에 여러 개의 모델을 사용할 때 혹은 모델의 다른 부분들을 여러개의 optimizer를 사용할 때 유용하다고 보면 됩니다.

 

즉,

1) optimizer 1개 - model 여러개

2) model 1개 - optimizer 여러개

 

이런 상황에서는 각 상황에 맞게 zero_grad()를 설정함으로써 학습이 가능하게 끔 할 수 있습니다.

1)의 경우는 optimizer.zero_grad()가 올바른 방법일 것 같고 2)의 경우는 model.zero_grad()가 올바른 방법인 것 같습니다.

 

이와 관련한 stackoverflow에 좋은 설명이 있으니 참조하시길 바랍니다.

 

stackoverflow.com/questions/61898668/net-zero-grad-vs-optim-zero-grad-pytorch

net.zero_grad() vs optim.zero_grad() pytorch

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

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