torch.optim¶
torch.optim 是一個實作各種優化演算法的套件。
最常用的方法都已支援,而且介面足夠通用,因此未來也可以輕鬆整合更複雜的方法。
如何使用優化器¶
要使用 torch.optim,您必須建構一個優化器物件,該物件將保存目前的狀態,並根據計算出的梯度更新參數。
建構它¶
要建構 Optimizer,您必須給它一個包含要優化的參數的可迭代物件(所有參數都應該是 Variable)。然後,您可以指定特定於優化器的選項,例如學習率、權重衰減等。
範例
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)
每個參數的選項¶
Optimizer 也支援指定每個參數的選項。為此,請傳遞一個 dict 的可迭代物件,而不是傳遞一個 Variable 的可迭代物件。它們每一個都將定義一個單獨的參數群組,並且應該包含一個 params 鍵,其中包含屬於該群組的參數清單。其他鍵應該與優化器接受的關鍵字引數相符,並將用作該群組的優化選項。
例如,當您想要指定每層學習率時,這非常有用
optim.SGD([
                {'params': model.base.parameters(), 'lr': 1e-2},
                {'params': model.classifier.parameters()}
            ], lr=1e-3, momentum=0.9)
這表示 model.base 的參數將使用 1e-2 的學習率,而 model.classifier 的參數將堅持使用預設的 1e-3 的學習率。最後,所有參數都將使用 0.9 的動量。
注意
您仍然可以將選項作為關鍵字引數傳遞。它們將在未覆蓋它們的群組中用作預設值。當您只想改變一個選項,同時保持所有其他選項在參數群組之間一致時,這非常有用。
還要考慮以下與參數的不同懲罰相關的範例。請記住,parameters() 會返回一個可迭代物件,其中包含所有可學習的參數,包括偏差和其他可能需要不同懲罰的參數。為了處理這個問題,可以為每個參數群組指定個別的懲罰權重
bias_params = [p for name, p in self.named_parameters() if 'bias' in name]
others = [p for name, p in self.named_parameters() if 'bias' not in name]
optim.SGD([
                {'params': others},
                {'params': bias_params, 'weight_decay': 0}
            ], weight_decay=1e-2, lr=1e-2)
通過這種方式,偏差項與非偏差項隔離開來,並且專門為偏差項設置了 weight_decay 為 0,以避免對該群組進行任何懲罰。
執行優化步驟¶
所有優化器都實作了一個 step() 方法,用於更新參數。它可以通過兩種方式使用
optimizer.step()¶
這是大多數優化器都支援的簡化版本。該函數可以在使用例如 backward() 計算出梯度後調用。
範例
for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()
optimizer.step(closure)¶
一些優化演算法(例如共軛梯度和 LBFGS)需要多次重新評估函數,因此您必須傳遞一個閉包,允許它們重新計算您的模型。閉包應該清除梯度,計算損失,並返回它。
範例
for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        return loss
    optimizer.step(closure)
基類¶
- class torch.optim.Optimizer(params, defaults)[原始碼]¶
- 所有優化器的基類。 - 警告 - 需要將參數指定為具有確定性順序的集合,該順序在多次運行中保持一致。不滿足這些屬性的物件的例子包括集合和對字典的值進行迭代的迭代器。 - 參數
- params (iterable) – - torch.Tensor或- dict的可迭代物件。指定要優化的張量。
- defaults (Dict[str, Any]) – (dict): 一個包含優化選項預設值的字典(當參數群組未指定時使用)。 
 
 
| 將參數群組添加到  | |
| 載入優化器狀態。 | |
| 以  | |
| 執行單個優化步驟(參數更新)。 | |
| 重置所有優化過的  | 
演算法¶
| 實作 Adadelta 演算法。 | |
| 實作 Adagrad 演算法。 | |
| 實作 Adam 演算法。 | |
| 實作 AdamW 演算法。 | |
| SparseAdam 實作了適用於稀疏梯度的 Adam 演算法的遮罩版本。 | |
| 實作 Adamax 演算法(基於無窮範數的 Adam 變體)。 | |
| 實作平均隨機梯度下降。 | |
| 實作 L-BFGS 演算法。 | |
| 實作 NAdam 演算法。 | |
| 實作 RAdam 演算法。 | |
| 實作 RMSprop 演算法。 | |
| 實作彈性反向傳播演算法。 | |
| 實作隨機梯度下降(可選擇使用動量)。 | 
我們的許多演算法都有各種針對性能、可讀性和/或通用性進行了優化的實作,因此如果用戶沒有指定特定的實作,我們會嘗試默認使用當前設備上通常最快的實作。
我們有 3 個主要的實作類別:for 迴圈、foreach(多張量)和融合。最直接的實作是使用大量計算的參數 for 迴圈。for 迴圈通常比我們的 foreach 實作慢,後者將參數組合成一個多張量並一次性運行大量計算,從而節省了許多順序內核調用。我們的一些優化器甚至還有更快的融合實作,它們將大量計算融合到一個內核中。我們可以將 foreach 實作視為水平融合,將融合實作視為在此基礎上的垂直融合。
一般來說,3 種實作的性能順序是融合 > foreach > for 迴圈。因此,在適用的情況下,我們默認使用 foreach 而不是 for 迴圈。適用意味著 foreach 實作可用,用戶沒有指定任何特定於實作的 kwargs(例如,fused、foreach、differentiable),並且所有張量都是原生的並且在 CUDA 上。請注意,雖然融合應該比 foreach 更快,但這些實作更新,我們希望在全面切換之前給予它們更多的磨合時間。歡迎您試用它們!
下表顯示了每種演算法的可用和默認實作
| 演算法 | 默認 | 有 foreach 嗎? | 有融合嗎? | 
|---|---|---|---|
| foreach | 是 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 是 | |
| foreach | 是 | 是 | |
| for 迴圈 | 否 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 否 | |
| for 迴圈 | 否 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 否 | |
| foreach | 是 | 否 | 
如何調整學習率¶
torch.optim.lr_scheduler 提供了几種基於 epochs 數量調整學習率的方法。torch.optim.lr_scheduler.ReduceLROnPlateau 允許基於某些驗證度量動態降低學習率。
學習率調整應該在優化器更新之後應用;例如,您應該這樣編寫代碼
範例
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler = ExponentialLR(optimizer, gamma=0.9)
for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler.step()
大多數學習率排程器可以背靠背調用(也稱為鏈式排程器)。結果是每個排程器將在前一個排程器獲得的學習率的基礎上依次應用。
範例
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
scheduler1 = ExponentialLR(optimizer, gamma=0.9)
scheduler2 = MultiStepLR(optimizer, milestones=[30,80], gamma=0.1)
for epoch in range(20):
    for input, target in dataset:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
    scheduler1.step()
    scheduler2.step()
在文檔的許多地方,我們將使用以下模板來指代排程器演算法。
>>> scheduler = ...
>>> for epoch in range(100):
>>>     train(...)
>>>     validate(...)
>>>     scheduler.step()
警告
在 PyTorch 1.1.0 之前,學習率排程器應該在優化器更新之前調用;1.1.0 以不兼容舊版本的方式改變了這種行為。如果您在優化器更新(調用 optimizer.step())之前使用學習率排程器(調用 scheduler.step()),這將跳過學習率排程的第一個值。如果您在升級到 PyTorch 1.1.0 後無法重現結果,請檢查您是否在錯誤的時間調用了 scheduler.step()。
| 將每個參數組的學習率設置為初始 lr 乘以給定函數的值。 | |
| 將每個參數組的學習率乘以指定函數中給出的因子。 | |
| 每隔 step_size 個 epochs,將每個參數組的學習率降低 gamma 倍。 | |
| 一旦 epoch 數量達到其中一個里程碑,就將每個參數組的學習率降低 gamma 倍。 | |
| 將每個參數組的學習率乘以一個小的常數因子,直到 epoch 數量達到預定義的里程碑:total_iters。 | |
| 通過線性改變小的乘法因子來降低每個參數組的學習率,直到 epoch 數量達到預定義的里程碑:total_iters。 | |
| 每個 epoch 都將每個參數組的學習率降低 gamma 倍。 | |
| 使用給定 total_iters 中的多項式函數來降低每個參數組的學習率。 | |
| 使用餘弦退火時間表設置每個參數組的學習率,其中 設置為初始 lr, 是自 SGDR 中上次重新啟動以來的 epoch 數 | |
| 鏈接學習率排程器列表。 | |
| 接收預計在優化過程中按順序調用的排程器列表,以及提供精確間隔的里程碑點,以反映應該在給定 epoch 調用哪個排程器。 | |
| 當指標停止改進時降低學習率。 | |
| 根據循環學習率策略 (CLR) 設置每個參數組的學習率。 | |
| 根據 1cycle 學習率策略設置每個參數組的學習率。 | |
| 使用餘弦退火時間表設置每個參數組的學習率,其中 設置為初始 lr, 是自上次重新啟動以來的 epoch 數, 是 SGDR 中兩次熱重新啟動之間的 epoch 數 | 
權重平均(SWA 和 EMA)¶
torch.optim.swa_utils 實作了隨機權重平均 (SWA) 和指數移動平均 (EMA)。特別是,torch.optim.swa_utils.AveragedModel 類別實作了 SWA 和 EMA 模型,torch.optim.swa_utils.SWALR 實作了 SWA 學習率排程器,而 torch.optim.swa_utils.update_bn() 是一個在訓練結束時用於更新 SWA/EMA 批次標準化統計量的工具函數。
SWA 已在 平均權重導致更廣泛的最佳值和更好的泛化 中提出。
EMA 是一種廣為人知的技術,透過減少所需的權重更新次數來縮短訓練時間。它是 Polyak 平均 的一種變體,但使用指數權重而不是跨迭代的相等權重。
構建平均模型¶
AveragedModel 類別用於計算 SWA 或 EMA 模型的權重。
您可以透過執行以下程式碼來建立 SWA 平均模型
>>> averaged_model = AveragedModel(model)
EMA 模型是透過指定 multi_avg_fn 參數來構建的,如下所示
>>> decay = 0.999
>>> averaged_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(decay))
衰減是一個介於 0 到 1 之間的參數,用於控制平均參數衰減的速度。如果未提供給 get_ema_multi_avg_fn,則預設值為 0.999。
get_ema_multi_avg_fn 會傳回一個函數,該函數會將以下 EMA 方程式套用至權重
其中 alpha 是 EMA 衰減率。
這裡的模型 model 可以是任意的 torch.nn.Module 物件。averaged_model 將會持續追蹤 model 參數的運行平均值。要更新這些平均值,您應該在 optimizer.step() 之後使用 update_parameters() 函數
>>> averaged_model.update_parameters(model)
對於 SWA 和 EMA,此呼叫通常在優化器 step() 之後立即執行。在 SWA 的情況下,這通常會在訓練開始時跳過一些步驟。
自訂平均策略¶
預設情況下,torch.optim.swa_utils.AveragedModel 會計算您提供的參數的運行相等平均值,但您也可以使用自訂平均函數,透過 avg_fn 或 multi_avg_fn 參數
- avg_fn允許定義一個對每個參數元組(平均參數、模型參數)進行操作的函數,並且應該傳回新的平均參數。
- multi_avg_fn允許定義對參數列表元組(平均參數列表、模型參數列表)同時進行更有效率的操作,例如使用- torch._foreach*函數。此函數必須就地更新平均參數。
在以下範例中,ema_model 使用 avg_fn 參數計算指數移動平均值
>>> ema_avg = lambda averaged_model_parameter, model_parameter, num_averaged:\
>>>         0.9 * averaged_model_parameter + 0.1 * model_parameter
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, avg_fn=ema_avg)
在以下範例中,ema_model 使用更有效率的 multi_avg_fn 參數計算指數移動平均值
>>> ema_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(0.9))
SWA 學習率排程¶
通常,在 SWA 中,學習率會設定為一個較高的常數值。SWALR 是一個學習率排程器,它會將學習率退火到一個固定值,然後保持不變。例如,以下程式碼建立了一個排程器,它會在每個參數群組中,將學習率從其初始值線性退火到 0.05,持續 5 個時期
>>> swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, \
>>>         anneal_strategy="linear", anneal_epochs=5, swa_lr=0.05)
您也可以透過設定 anneal_strategy="cos",使用餘弦退火到固定值,而不是線性退火。
處理批次標準化¶
update_bn() 是一個工具函數,允許在訓練結束時,在給定的資料載入器 loader 上計算 SWA 模型的批次標準化統計量
>>> torch.optim.swa_utils.update_bn(loader, swa_model)
update_bn() 會將 swa_model 套用至資料載入器中的每個元素,並計算模型中每個批次標準化層的激活統計量。
警告
update_bn() 假設資料載入器 loader 中的每個批次都是張量或張量列表,其中第一個元素是網路 swa_model 應該套用的張量。如果您的資料載入器具有不同的結構,您可以透過使用 swa_model 對資料集的每個元素進行正向傳遞,來更新 swa_model 的批次標準化統計量。
整合所有內容:SWA¶
在以下範例中,swa_model 是累積權重平均值的 SWA 模型。我們訓練模型總共 300 個時期,並在第 160 個時期切換到 SWA 學習率排程,並開始收集參數的 SWA 平均值
>>> loader, optimizer, model, loss_fn = ...
>>> swa_model = torch.optim.swa_utils.AveragedModel(model)
>>> scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=300)
>>> swa_start = 160
>>> swa_scheduler = SWALR(optimizer, swa_lr=0.05)
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>       if epoch > swa_start:
>>>           swa_model.update_parameters(model)
>>>           swa_scheduler.step()
>>>       else:
>>>           scheduler.step()
>>>
>>> # Update bn statistics for the swa_model at the end
>>> torch.optim.swa_utils.update_bn(loader, swa_model)
>>> # Use swa_model to make predictions on test data
>>> preds = swa_model(test_input)
整合所有內容:EMA¶
在以下範例中,ema_model 是 EMA 模型,它以 0.999 的衰減率累積權重的指數衰減平均值。我們訓練模型總共 300 個時期,並立即開始收集 EMA 平均值。
>>> loader, optimizer, model, loss_fn = ...
>>> ema_model = torch.optim.swa_utils.AveragedModel(model, \
>>>             multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))
>>>
>>> for epoch in range(300):
>>>       for input, target in loader:
>>>           optimizer.zero_grad()
>>>           loss_fn(model(input), target).backward()
>>>           optimizer.step()
>>>           ema_model.update_parameters(model)
>>>
>>> # Update bn statistics for the ema_model at the end
>>> torch.optim.swa_utils.update_bn(loader, ema_model)
>>> # Use ema_model to make predictions on test data
>>> preds = ema_model(test_input)