捷徑

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_decay0,以避免對該群組進行任何懲罰。

執行優化步驟

所有優化器都實作了一個 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.Tensordict 的可迭代物件。指定要優化的張量。

  • defaults (Dict[str, Any]) – (dict): 一個包含優化選項預設值的字典(當參數群組未指定時使用)。

Optimizer.add_param_group

將參數群組添加到 Optimizerparam_groups

Optimizer.load_state_dict

載入優化器狀態。

Optimizer.state_dict

dict 的形式返回優化器的狀態。

Optimizer.step

執行單個優化步驟(參數更新)。

Optimizer.zero_grad

重置所有優化過的 torch.Tensor 的梯度。

演算法

Adadelta

實作 Adadelta 演算法。

Adagrad

實作 Adagrad 演算法。

Adam

實作 Adam 演算法。

AdamW

實作 AdamW 演算法。

SparseAdam

SparseAdam 實作了適用於稀疏梯度的 Adam 演算法的遮罩版本。

Adamax

實作 Adamax 演算法(基於無窮範數的 Adam 變體)。

ASGD

實作平均隨機梯度下降。

LBFGS

實作 L-BFGS 演算法。

NAdam

實作 NAdam 演算法。

RAdam

實作 RAdam 演算法。

RMSprop

實作 RMSprop 演算法。

Rprop

實作彈性反向傳播演算法。

SGD

實作隨機梯度下降(可選擇使用動量)。

我們的許多演算法都有各種針對性能、可讀性和/或通用性進行了優化的實作,因此如果用戶沒有指定特定的實作,我們會嘗試默認使用當前設備上通常最快的實作。

我們有 3 個主要的實作類別:for 迴圈、foreach(多張量)和融合。最直接的實作是使用大量計算的參數 for 迴圈。for 迴圈通常比我們的 foreach 實作慢,後者將參數組合成一個多張量並一次性運行大量計算,從而節省了許多順序內核調用。我們的一些優化器甚至還有更快的融合實作,它們將大量計算融合到一個內核中。我們可以將 foreach 實作視為水平融合,將融合實作視為在此基礎上的垂直融合。

一般來說,3 種實作的性能順序是融合 > foreach > for 迴圈。因此,在適用的情況下,我們默認使用 foreach 而不是 for 迴圈。適用意味著 foreach 實作可用,用戶沒有指定任何特定於實作的 kwargs(例如,fused、foreach、differentiable),並且所有張量都是原生的並且在 CUDA 上。請注意,雖然融合應該比 foreach 更快,但這些實作更新,我們希望在全面切換之前給予它們更多的磨合時間。歡迎您試用它們!

下表顯示了每種演算法的可用和默認實作

演算法

默認

有 foreach 嗎?

有融合嗎?

Adadelta

foreach

Adagrad

foreach

Adam

foreach

AdamW

foreach

SparseAdam

for 迴圈

Adamax

foreach

ASGD

foreach

LBFGS

for 迴圈

NAdam

foreach

RAdam

foreach

RMSprop

foreach

Rprop

foreach

SGD

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_scheduler.LambdaLR

將每個參數組的學習率設置為初始 lr 乘以給定函數的值。

lr_scheduler.MultiplicativeLR

將每個參數組的學習率乘以指定函數中給出的因子。

lr_scheduler.StepLR

每隔 step_size 個 epochs,將每個參數組的學習率降低 gamma 倍。

lr_scheduler.MultiStepLR

一旦 epoch 數量達到其中一個里程碑,就將每個參數組的學習率降低 gamma 倍。

lr_scheduler.ConstantLR

將每個參數組的學習率乘以一個小的常數因子,直到 epoch 數量達到預定義的里程碑:total_iters。

lr_scheduler.LinearLR

通過線性改變小的乘法因子來降低每個參數組的學習率,直到 epoch 數量達到預定義的里程碑:total_iters。

lr_scheduler.ExponentialLR

每個 epoch 都將每個參數組的學習率降低 gamma 倍。

lr_scheduler.PolynomialLR

使用給定 total_iters 中的多項式函數來降低每個參數組的學習率。

lr_scheduler.CosineAnnealingLR

使用餘弦退火時間表設置每個參數組的學習率,其中 ηmax\eta_{max} 設置為初始 lr,TcurT_{cur} 是自 SGDR 中上次重新啟動以來的 epoch 數

lr_scheduler.ChainedScheduler

鏈接學習率排程器列表。

lr_scheduler.SequentialLR

接收預計在優化過程中按順序調用的排程器列表,以及提供精確間隔的里程碑點,以反映應該在給定 epoch 調用哪個排程器。

lr_scheduler.ReduceLROnPlateau

當指標停止改進時降低學習率。

lr_scheduler.CyclicLR

根據循環學習率策略 (CLR) 設置每個參數組的學習率。

lr_scheduler.OneCycleLR

根據 1cycle 學習率策略設置每個參數組的學習率。

lr_scheduler.CosineAnnealingWarmRestarts

使用餘弦退火時間表設置每個參數組的學習率,其中 ηmax\eta_{max} 設置為初始 lr,TcurT_{cur} 是自上次重新啟動以來的 epoch 數,TiT_{i} 是 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 方程式套用至權重

Wt+1EMA=αWtEMA+(1α)WtmodelW^\textrm{EMA}_{t+1} = \alpha W^\textrm{EMA}_{t} + (1 - \alpha) W^\textrm{model}_t

其中 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_fnmulti_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)

文件

取得 PyTorch 的完整開發人員文件

檢視文件

教學課程

取得適用於初學者和進階開發人員的深入教學課程

檢視教學課程

資源

尋找開發資源並獲得問題解答

檢視資源