[기본] optimizer.zero_grad()는 어디에 넣어주어야 할까?

모델을 훈련시킬 때 사람에 따라 optimizer.zero_grad()가 맨 처음에 있기도 하고, 아래 코드처럼 맨 뒤에 있기도 하다. 어디에 넣는 게 정석코드인지 신경쓰였는데, 결론적으로는 loss.backward()와 optimizer.step() 사이에 넣지만 않으면 상관없다고 한다. 애초에 이 명령어가 모델을 학습시킬 때 loss.backward()로 변화하는 기울기값 (loss.grad) 들이 계속 누적되지 않도록 0으로 초기화해주는 역할을 하기 때문이다.

* RNN 등의 모듈에서는 기울기값이 계속 누적되는 것이 계산에 있어 편리하기 때문에 pytorch에서는 기본적으로 gradient accumulation을 하고 있다

#pseudo-code

for x in train_batch:
	...
	loss, logits, _, _ = model(**x)
    loss.backward()
    optimizer.step()
    scheduler.step()
    optimizer.zero_grad()