圖解OneFlow的學習率調整策略

語言: CN / TW / HK

 

撰文|李佳

 

1

背景

 

學習率調整策略(learning rate scheduler),其實單獨拎出每一個來看都不難,但是由於方法較多,上來就看文件容易一頭霧水, 以OneFlow v0.7.0為例,oneflow.optim.lr_scheduler模組中就包含了14種策略。

 

有沒有一種更好的方法來學習呢?比如可視化出學習率的變化過程,此時,我腦海中突然浮現出Convolution Arithmetic這個經典專案,作者將各種CNN卷積操作以gif形式展示,一目瞭然。

 

 

所以,就有了這篇文章,將學習率調整策略可視化出來,下面是兩個例子(ConstantLR和LinearLR):

 

 

 

我將視覺化程式碼分別託管在Hugging Face Spaces和Streamlit Cloud,大家可以任選一個連結訪問,然後自由調節引數,感受學習率的變化過程。

 

  • https://huggingface.co/spaces/basicv8vc/learning-rate-scheduler-online

  • https://share.streamlit.io/basicv8vc/scheduler-online

 

2

學習率調整策略

 

學習率可以說是訓練神經網路過程中最重要的引數(之一),目前大家都已接受用動態學習率調整策略來代替固定學習率,各種學習率調整策略層出不窮,下面我們就以OneFlow v0.7.0為例,學習下常用的幾種策略。

 

基類LRScheduler

 

LRScheduler(optimizer: Optimizer, last_step: int = -1, verbose: bool = False)是所有學習率排程器的基類,初始化引數中last_step和verbose一般不需要設定,前者主要和checkpoint相關,後者則是在每次step() 呼叫時列印學習率,可以用於 debug。LRScheduler中最重要的方法是step(),這個方法的作用就是修改使用者設定的初始學習率,然後應用到下一次的Optimizer.step()。

 

有些資料會講LRScheduler根據epoch或iteration/step來調整學習率,兩種說法都沒問題,實際上,LRScheduler並不知道當前訓練到第幾個epoch或第幾個iteration/step,只記錄了呼叫step()的次數(last_step),如果每個epoch呼叫一次,那就是根據epoch來調整學習率,如果每個mini-batch呼叫一次,那就是根據iteration來調整學習率。以訓練Transformer模型為例,需要在每個iteration呼叫step()。

 

簡單來說,LRScheduler根據調整策略本身、當前呼叫step()的次數(last_step)和使用者設定的初始學習率來得到下一次梯度更新時的學習率。

 

ConstantLR

 

oneflow.optim.lr_scheduler.ConstantLR(    optimizer: Optimizer,    factor: float = 1.0 / 3,    total_iters: int = 5,    last_step: int = -1,    verbose: bool = False,)

 

ConstantLR和固定學習率差不多,唯一的區別是在前total_iters,學習率為初始學習率 * factor。

 

注意:由於factor取值[0, 1],所以這是一個學習率遞增的策略。

 

ConstantLR

 

LinearLR

 

oneflow.optim.lr_scheduler.LinearLR(    optimizer: Optimizer,    start_factor: float = 1.0 / 3,    end_factor: float = 1.0,    total_iters: int = 5,    last_step: int = -1,    verbose: bool = False,)

 

LinearLR和固定學習率也差不多,唯一的區別是在前total_iters,學習率先線性增加或遞減,然後再固定為初始學習率 * end_factor。

 

 

注意:學習率在前total_iters是遞增or遞減由start_factor和end_factor大小決定。

 

LinearLR

 

ExponentialLR

 

oneflow.optim.lr_scheduler.ExponentialLR(    optimizer: Optimizer,    gamma: float,    last_step: int = -1,    verbose: bool = False,)

 

學習率呈指數衰減,當然也可以將gamma設定為>1,進行指數增加,不過估計沒人願意這麼做。

 

 

ExponentialLR

 

StepLR

 

oneflow.optim.lr_scheduler.StepLR(    optimizer: Optimizer,    step_size: int,    gamma: float = 0.1,    last_step: int = -1,    verbose: bool = False,)

 

StepLR和ExponentialLR差不多,區別是不是每一次呼叫step()都進行學習率調整,而是每隔step_size才調整一次。

 

StepLR

 

MultiStepLR

 

oneflow.optim.lr_scheduler.MultiStepLR(    optimizer: Optimizer,    milestones: list,    gamma: float = 0.1,    last_step: int = -1,    verbose: bool = False,)

 

StepLR每隔step_size就調整一次學習率,而MultiStepLR則根據使用者指定的milestones進行調整,假設milestones是[2, 5, 9],在[0, 2)是lr,在[2, 5)是lr * gamma,在[5, 9)是lr * (gamma **2),在[9, )是lr * (gamma **3)。

 

MultiStepLR

 

PolynomialLR

 

oneflow.optim.lr_scheduler.PolynomialLR(    optimizer,    steps: int,    end_learning_rate: float = 0.0001,    power: float = 1.0,    cycle: bool = False,    last_step: int = -1,    verbose: bool = False,)

 

前面的學習率調整策略無非是線性或指數,PolynomialLR則根據多項式進行調整,先看cycle引數,預設是False,此時先根據多項式衰減然後再固定學習率,公式如下:

 

注:公式中的decay_batch就是steps,current_batch就是最新的last_step。

 

如果cycle是True,則稍微複雜點,類似於以steps為週期進行變化,每次從一個最大學習率衰減到end_learning_rate,每個週期的最大學習率也是逐漸衰減的,公式如下:

PolynomialLR

 

看下cycle=True的例子,

 

 

CosineDecayLR

 

oneflow.optim.lr_scheduler.CosineDecayLR(    optimizer: Optimizer,    decay_steps: int,    alpha: float = 0.0,    last_step: int = -1,    verbose: bool = False,)

 

在前decay_steps步,學習率由lr餘弦衰減到lr * alpha,然後固定為lr*alpha。

 

注:CosineDecayLR是為了對齊TensorFlow中的CosineDecay。

 

 

 

CosineAnnealingLR

 

oneflow.optim.lr_scheduler.CosineAnnealingLR(    optimizer: Optimizer,    T_max: int,    eta_min: float = 0.0,    last_step: int = -1,    verbose: bool = False,)

 

CosineAnnealingLR和CosineDecayLR很像,區別在於前者不僅包含餘弦衰減的過程,也可以包含餘弦增加,在前T_max步,學習率由lr餘弦衰減到eta_min, 如果cur_step > T_max,然後再餘弦增加到lr,不斷重複這個過程。

 

CosineAnnealingLR

 

CosineAnnealingWarmRestarts

 

oneflow.optim.lr_scheduler.CosineAnnealingWarmRestarts(    optimizer: Optimizer,    T_0: int,    T_mult: int = 1,    eta_min: float = 0.0,    decay_rate: float = 1.0,    restart_limit: int = 0,    last_step: int = -1,    verbose: bool = False,)

 

上面三個Cosine相關的LRScheduler來自同一篇論文(SGDR: Stochastic Gradient Descent with Warm Restarts),這個引數比較多,首先看T_mul,如果T_mul=1,則學習率等週期變化,週期大小就是T_0,也就是由最大學習率衰減到最小學習率的步數(steps),注意如果decay_rate<1,則每個週期的最大學習率和最小學習率都在衰減,第一個週期由lr開始衰減,第二個週期由lr * decay_rate開始衰減,第三個週期由lr * (decay_rate ** 2)開始衰減。

 

如果T_mult>1,則學習率不是等週期變化,每個週期的大小是上一個週期大小T_mult,第一個週期是T_0,第二個週期是T_0 * T_mult,第三個週期是 T_0 * T_mult * T_mult。

 

再來看restart_limit,預設值是0,就是上面的過程,如果>0,物理含義是週期數量,假設為3,則只有三次從最大衰減到最小,然後學習率一直是eta_min,不再週期變化了。

 

先看個T_mult=1的例子,此時decay_rate=1,

 

T_mult=1, decay_rate=1

 

再看個T_mult=1,decay_rate=0.5的例子,注意這種組合形式並不常用。

 

T_mult=1, decay_rate=0.5

 

再來看T_mult >1的例子,

 

 

最後,再看個restart_limit != 0的例子,

 

 

 

3

組合排程策略

 

上面講的都是單個學習率排程策略,再來看幾個學習率組合排程策略,比如訓練Transformer常用的Noam scheduler就需要先線性增加再指數衰減,可以通過LinearLR和ExponentialLR組合得到。也可以直接使用LambdaLR傳入學習率變化函式。

 

LambdaLR

oneflow.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_step=-1, verbose=False)

 

LambdaLR可以說是最靈活的策略了,因為具體的方法是根據函式lr_lambda來指定的。比如實現Transformer中的Noam Scheduler:

 

def rate(step, model_size, factor, warmup):    """    we have to default the step to 1 for LambdaLR function    to avoid zero raising to negative power.    """    if step == 0:        step = 1    return factor * (        model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))    )

model = CustomTransformer(...)optimizer = flow.optim.Adam(    model.parameters(), lr=1.0, betas=(0.9, 0.98), eps=1e-9)lr_scheduler = LambdaLR(    optimizer=optimizer,    lr_lambda=lambda step: rate(step, d_model, factor=1, warmup=3000),)

 

注意:OneFlow的Graph模式並不支援LambdaLR。

 

SequentialLR

 

oneflow.optim.lr_scheduler.SequentialLR(    optimizer: Optimizer,    schedulers: Sequence[LRScheduler],    milestones: Sequence[int],    interval_rescaling: Union[Sequence[bool], bool] = False,    last_step: int = -1,    verbose: bool = False,)

 

支援傳入多個LRScheduler,每個LRScheduler的作用範圍(step range)由milestones指定,主要看下interval_rescaling這個引數,預設是False,目的是讓相鄰的兩個scheduler在銜接時學習率比較平滑,比如milestones=[5],當last_step=5時,第二個schduler就從last_step=5開始計算新的學習率,這樣和last_step=4(前一個scheduler計算學習率)得到的學習率不會有過大差異,而interval_rescaling=True時,則這個scheduler的last_step從0開始。

 

WarmupLR

 

oneflow.optim.lr_scheduler.WarmupLR(    scheduler_or_optimizer: Union[LRScheduler, Optimizer],    warmup_factor: float = 1.0 / 3,    warmup_iters: int = 5,    warmup_method: str = "linear",    warmup_prefix: bool = False,    last_step=-1,    verbose=False,)

 

WarmupLR是SequentialLR的子類,包含兩個LRScheduler,並且第一個要麼是ConstantLR,要麼是LinearLR。

 

ChainedScheduler

 

oneflow.optim.lr_scheduler.ChainedScheduler(schedulers)

 

前面講的組合形式的排程策略,在每一個step,只有一個LRScheduler發揮作用,而ChainedScheduler,在每一個step計算學習率時,所有的LRScheduler都參與,類似於管道(pipeline)

 

lr ==> LRScheduler_1 ==> LRScheduler_2 ==> ... ==> LRScheduler_N

 

ReduceLROnPlateau

 

oneflow.optim.lr_scheduler.ReduceLROnPlateau(    optimizer,    mode="min",    factor=0.1,    patience=10,    threshold=1e-4,    threshold_mode="rel",    cooldown=0,    min_lr=0,    eps=1e-8,    verbose=False,)

 

前面提到的所有LRScheduler都是根據當前的step來計算學習率,而在模型訓練過程中,我們最關心的是訓練集和驗證集上面的指標,能不能利用這些指標來指導學習率變化呢?這時候可以用ReduceLROnPlateau,如果某項指標多個step都未發生顯著變化,則學習率進行線性衰減。

 

optimizer = flow.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)scheduler = flow.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')for epoch in range(10):    train(...)    val_loss = validate(...)    # 注意,該步驟應在validate()之後呼叫。    scheduler.step(val_loss)

 

4

實踐

 

如果看到這裡有點意猶未盡的感覺,不如動手實踐一下,下面是我根據官方的圖片分類例項改寫的CIFAR-100例子,可以設定不同的學習率排程策略來感受下差異

 

  • https://github.com/basicv8vc/oneflow-cifar100-lr-scheduler

 

(本文經授權後釋出,原文:

https://zhuanlan.zhihu.com/p/520719314 )

 

其他人都在看

歡迎體驗OneFlow v0.7.0:https://github.com/Oneflow-Inc/oneflow

 


本文分享自微信公眾號 - OneFlow(OneFlowTechnology)。
如有侵權,請聯絡 [email protected] 刪除。
本文參與“OSC源創計劃”,歡迎正在閱讀的你也加入,一起分享。