圖解OneFlow的學習率調整策略

 

撰文｜李佳

 

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 ）

 

其他人都在看

• 深度學習概述

• 一個運算元在深度學習框架中的旅程

• 手把手推導分散式矩陣乘的最優並行策略

• 訓練千億引數大模型，離不開四種並行策略

• 解讀Pathways（二）：向前一步是OneFlow

• 關於併發和並行，Go和Erlang之父都弄錯了？

• OneFlow v0.7.0釋出：全新分散式介面，LiBai、Serving等一應俱全

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

 

---