[Pytorch] Scaler 사용 시 Detected call of lr_scheduler.step() before optimizer.step() warning이 뜰 때

훈련을 할 때 optimizer 다음에 learning rate scheduler을 콜하는 게 자연스럽고, 실제로 그렇게 하지 않으면 맨 위의 사진처럼 경고가 뜨기도 한다. 하지만 scaler를 사용하면, 아래와 같이 맞게 적어도 초기 스텝에서 가끔 위의 에러가 뜰 때가 있다. 이는 scaler가 NaN, inf 같은 맞지 않는 미분값을 발견하면 optimizer.step() 단계를 건너뛰는데, scheduler는 이 사실을 모르기 때문이다.

#fp16으로 훈련
scaler = torch.cuda.amp.GradScaler()
for batch in train_ds:

	(...)
    
	scaler.step(optimizer)
	scheduler.step()
    scaler.update()
    

실제로 캡처한 사진도 scaler를 사용하고 있는데, 찾아보니 scaler.step(optimizer)가 반환하는 값이 없어서 optimizer.step()이 실행될 때만 scheduler.step()을 진행한다, 이런 식으로 코드를 짤 수가 없었다. 아래는 관련 의문을 담은 디스커션.

디스커션: https://discuss.pytorch.org/t/userwarning-detected-call-of-lr-scheduler-step-before-optimizer-step-in-pytorch-1-1-0-and-later-you-should-call-them-in-the-opposite-order-optimizer-step-before-lr-scheduler-step/88295
scaler 관련 torch site: https://tutorials.pytorch.kr/recipes/recipes/amp_recipe.html

새로운 기능을 만들겠다! 라고 20년에 답변이 되었지만 21년 9월에 아직도 안 만들어졌냐는 질문이 있고, 현재까지도 없다...

현재로서는 직접 gradient 가 줄어들었는지 비교하는 코드를 짤 게 아니라면, 초기 경고는 무시하는 게 마음 편하겠다. 실제로 맨 위 사진을 보면 저 경고가 전체 코드에서 총 2번 나온다.