pytorch-lr-finder

Intro

평소에 pytroch_forecasting Library를 쓰면 보통 pytorch_lightning과 결합하여서 써서, auto_lr_find를 이용해서 최적의 learning rate를 찾곤했다. 하지만 내가 직접 모델을 구현할때는 보통 pytorch_lightning을 사용하지 않고 손수 loss를 구한후 backward를 시키면서 모델을 훈련시키는데, 이 과정에서 초기 learning rate를 heuristic하게 찾아야 해서 좀 불편했다. 이러한 니즈를 해결하고자 검색해본결과
GitHub
에 올라와있는 모듈의 README를 한번 공부해보고 실제로 사용해보겠다.

이 라이브러리는 Paper에 자세히 기술된 learning rate range test(LRRT)의 파이토치 구현이다.

LRRT란?
learning rate range test는 최적 learning rate에 대한 가치있는 정보를 제공하는 테스트이다. 작은 learning rate로 시작하면 network가 수렴하기 시작하는데, 너무 작다면 수렴하지 않고 멈추는 현상이 발생할 것이다. 이러한 Trouble을 없애기 위해서 learning rate를 키우다보면, 결국에는 network가 발산해버릴 것이다.

Typically하게 좋은 Static(정적인, Scheduler로 조정하지 않는) learning rate는 감소하는 loss curve에서 half정도의 수준일 수 있다. 아래의 plot에서는 0.002가 좋다.

위에 Paper에서 기술된 cyclical learning rate란 어떤 영역사이(start_lr,end_lr)에서 learning rate가 순회하는 것이다.
저자는 loss가 감소하기 시작하는 부분을 start_lr 더이상 감소하지 않는 부분을 end_lr로 설정하도록 조언한다.
위의 그림에서는 start_lr =0.0002 and end_lr=0.2이다.

Implementation details and usage

fastai로부터 수정된 version

learning rate를 지수방식으로 증가시키고 train loss를 각 learning rate에 대해 계산한다.

from torch_lr_finder import LRFinder

model = ...
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-7, weight_decay=1e-2)
lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, end_lr=100, num_iter=100)
lr_finder.plot() # to inspect the loss-learning rate graph
lr_finder.reset() # to reset the model and optimizer to their initial state

Leslie Smith's의 접근

각각의 learning rate에 대해서 learning rate를 선형적으로 증가시키고 evaluation loss를 계산한다.
lr_finder.plot()은 learning rate에 따른 evaluation loss를 plot해준다. 이러한 Leslie의 접근은 evaluation loss를 사용한다는 점에서 더욱 발산하기 쉽기 때문에 더 정확한 그래프를 생산한다.

하지만 이러한 접근은 시간적으로 test를 수행하는데 오래걸리고 특히나 evaluation dataset이 크다면 더욱 오래 걸린다는 단점이 있다.
아래의 코드를 확인해보면 trainloaderd와 동시에 val_loader를 제공한 것을 볼 수 있다.

from torch_lr_finder import LRFinder

model = ...
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.1, weight_decay=1e-2)
lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
lr_finder.range_test(trainloader, val_loader=val_loader, end_lr=1, num_iter=100, step_mode="linear")
lr_finder.plot(log_lr=False)
lr_finder.reset()

NOTE

• LRFinder에 passing된 optimizer는 LRScheduler를 가져서는 안된다.
• LRFinder.range_test()는 모델의 weight와 optimizer의 parameter들을 변화시키는데, 이는 LRFinder.reset()으로 원래대로 복구 가능하다.
• learning rate와 loss history는 lr_finder.history를 통해서 접근 가능하고 각각 lr과 loss key를 사용하면 된다.