快捷方式

SparseAdam

class torch.optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, maximize=False)[原始碼][原始碼]

SparseAdam 實現了適用於稀疏梯度的 Adam 演算法的掩碼版本。目前,由於實現限制(如下所述),SparseAdam 僅適用於一部分狹窄的用例,特別是具有稀疏佈局梯度的密集佈局引數。這發生在模組反向傳播已產生稀疏佈局梯度的一種特殊情況下。一個示例如此行為的神經網路模組是 nn.Embedding(sparse=True)

SparseAdam 透過遮蔽掉梯度中零值對應的引數和動量更新來近似 Adam 演算法。Adam 演算法會根據梯度的所有值更新一階動量、二階動量和引數,而 SparseAdam 只更新梯度中非零值對應的動量和引數。

一種簡化的、關於其 intended 實現方式的思考方法如下:

  1. 建立一個稀疏梯度中非零值的掩碼。例如,如果你的梯度看起來像 [0, 5, 0, 0, 9],那麼掩碼將是 [0, 1, 0, 0, 1]。

  2. 將此掩碼應用於執行動量,並且僅對非零值進行計算。

  3. 將此掩碼應用於引數,並且僅對非零值應用更新。

實際上,我們使用稀疏佈局張量來最佳化此近似方法,這意味著未具體化的梯度被掩蓋得越多,最佳化效能越好。由於我們依賴於使用稀疏佈局張量,我們推斷稀疏佈局中任何具體化的值都是非零的,並且我們實際上不會驗證所有值是否都不為零!重要的是不要混淆語義稀疏張量(其許多值為零的張量)與稀疏佈局張量(其 .is_sparse 返回 True 的張量)。SparseAdam 近似方法旨在用於語義稀疏張量,而稀疏佈局僅是實現細節。更清晰的實現方法是使用 MaskedTensors,但這些是實驗性的。

注意

如果你懷疑你的梯度在語義上是稀疏的(但不是稀疏佈局),則此變體可能不適合你。理想情況下,你首先應避免具體化任何被懷疑是稀疏的內容,因為將所有梯度從密集佈局轉換為稀疏佈局的開銷可能抵消效能提升。在這種情況下,使用 Adam 可能是最佳選擇,除非你可以輕鬆地調整模組使其輸出稀疏梯度,類似於 nn.Embedding(sparse=True)。如果你堅持轉換梯度,可以在呼叫 .step() 之前手動用稀疏等效項覆蓋引數的 .grad 欄位來完成此操作。

引數
  • params (可迭代物件) – 要最佳化的引數或 named_parameters 的可迭代物件,或定義引數組的字典的可迭代物件。使用 named_parameters 時,所有組中的所有引數都應命名

  • lr (浮點數, 張量, 可選) – 學習率 (預設值: 1e-3)

  • betas (元組[浮點數, 浮點數], 可選) – 用於計算梯度及其平方的執行平均值的係數 (預設值: (0.9, 0.999))

  • eps (浮點數, 可選) – 新增到分母以提高數值穩定性的項 (預設值: 1e-8)

  • maximize (布林值, 可選) – 相對於引數最大化目標,而不是最小化 (預設值: False)

add_param_group(param_group)[原始碼]

將一個引數組新增到 Optimizerparam_groups 中。

這在微調預訓練網路時很有用,因為可以在訓練過程中使凍結層可訓練並新增到 Optimizer 中。

引數

param_group (字典) – 指定應最佳化哪些 張量 以及組特定的最佳化選項。

load_state_dict(state_dict)[原始碼]

載入最佳化器狀態。

引數

state_dict (字典) – 最佳化器狀態。應是呼叫 state_dict() 返回的物件。

注意

引數的名稱(如果它們存在於 state_dict() 中每個引數組的“param_names”鍵下)不會影響載入過程。對於自定義用例(例如,載入的狀態字典中的引數與最佳化器中初始化的引數不同),應實現自定義的 register_load_state_dict_pre_hook 以相應地調整載入的字典。如果載入的狀態字典 param_groups 中存在 param_names,它們將被儲存並覆蓋最佳化器狀態中當前存在的名稱(如果存在)。如果它們不存在於載入的狀態字典中,最佳化器的 param_names 將保持不變。

register_load_state_dict_post_hook(hook, prepend=False)[原始碼]

註冊一個 load_state_dict 後置鉤子,該鉤子將在 load_state_dict() 呼叫後被呼叫。它應具有以下簽名

hook(optimizer) -> None

optimizer 引數是正在使用的最佳化器例項。

鉤子將在對 self 呼叫 load_state_dict 後以 self 為引數被呼叫。註冊的鉤子可用於在 load_state_dict 載入 state_dict 後執行後處理。

引數
  • hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

  • prepend (布林值) – 如果為 True,提供的後置鉤子將在所有已註冊的 load_state_dict 後置鉤子之前觸發。否則,提供的鉤子將在所有已註冊的後置鉤子之後觸發。(預設值: False)

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_load_state_dict_pre_hook(hook, prepend=False)[原始碼]

註冊一個 load_state_dict 前置鉤子,該鉤子將在 load_state_dict() 呼叫前被呼叫。它應具有以下簽名

hook(optimizer, state_dict) -> state_dict or None

optimizer 引數是正在使用的最佳化器例項,state_dict 引數是使用者傳入 load_state_dict 的 state_dict 的淺複製。鉤子可以就地修改 state_dict 或選擇返回一個新的。如果返回一個 state_dict,它將被用於載入到最佳化器中。

鉤子將在對 self 呼叫 load_state_dict 之前以 self 和 state_dict 為引數被呼叫。註冊的鉤子可用於在進行 load_state_dict 呼叫之前執行預處理。

引數
  • hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

  • prepend (布林值) – 如果為 True,提供的前置鉤子將在所有已註冊的 load_state_dict 前置鉤子之前觸發。否則,提供的前置鉤子將在所有已註冊的前置鉤子之後觸發。(預設值: False)

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_state_dict_post_hook(hook, prepend=False)[原始碼]

註冊一個 state dict 後置鉤子,該鉤子將在 state_dict() 呼叫後被呼叫。

它應具有以下簽名

hook(optimizer, state_dict) -> state_dict or None

鉤子將在對 self 生成 state_dict 後以 self 和 state_dict 為引數被呼叫。鉤子可以就地修改 state_dict 或選擇返回一個新的。註冊的鉤子可用於在 state_dict 返回之前對其進行後處理。

引數
  • hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

  • prepend (布林值) – 如果為 True,提供的後置鉤子將在所有已註冊的 state_dict 後置鉤子之前觸發。否則,提供的後置鉤子將在所有已註冊的後置鉤子之後觸發。(預設值: False)

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_state_dict_pre_hook(hook, prepend=False)[原始碼]

註冊一個 state dict 前置鉤子,該鉤子將在 state_dict() 呼叫前被呼叫。

它應具有以下簽名

hook(optimizer) -> None

optimizer 引數是正在使用的最佳化器例項。鉤子將在對 self 呼叫 state_dict 之前以 self 為引數被呼叫。註冊的鉤子可用於在進行 state_dict 呼叫之前執行預處理。

引數
  • hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

  • prepend (布林值) – 如果為 True,提供的前置鉤子將在所有已註冊的 state_dict 前置鉤子之前觸發。否則,提供的前置鉤子將在所有已註冊的前置鉤子之後觸發。(預設值: False)

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemoveableHandle

register_step_post_hook(hook)[原始碼]

註冊一個最佳化器步進後置鉤子,該鉤子將在最佳化器步進後被呼叫。

它應具有以下簽名

hook(optimizer, args, kwargs) -> None

optimizer 引數是正在使用的最佳化器例項。

引數

hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemovableHandle

register_step_pre_hook(hook)[原始碼]

註冊一個最佳化器步進前置鉤子,該鉤子將在最佳化器步進前被呼叫。

它應具有以下簽名

hook(optimizer, args, kwargs) -> None or modified args and kwargs

optimizer 引數是正在使用的最佳化器例項。如果 args 和 kwargs 被前置鉤子修改,則轉換後的值將作為包含 new_args 和 new_kwargs 的元組返回。

引數

hook (可呼叫物件) – 要註冊的使用者定義的鉤子。

返回值

一個控制代碼,可以透過呼叫 handle.remove() 來移除新增的鉤子

返回型別

torch.utils.hooks.RemovableHandle

state_dict()[原始碼]

將最佳化器的狀態以 dict 的形式返回。

它包含兩個條目

  • state: 一個 Dict,包含當前的最佳化狀態。其內容

    在不同的最佳化器類之間有所不同,但保留了一些共同特性。例如,狀態按引數儲存,但不儲存引數本身。 state 是一個將引數 ID 對映到包含與每個引數對應的狀態的 Dict 的字典。

  • param_groups: 一個 List,包含所有引數組,其中每個

    引數組是一個 Dict。每個引數組包含特定於最佳化器的元資料,例如學習率和權重衰減,以及該組中引數的引數 ID 列表。如果引數組是使用 named_parameters() 初始化的,則名稱內容也會儲存在狀態字典中。

注意:引數 ID 可能看起來像索引,但它們只是將狀態與 param_group 關聯起來的 ID。從 state_dict 載入時,最佳化器會將 param_group 的 params (整數 ID) 與最佳化器的 param_groups (實際的 nn.Parameter) 壓縮,以便匹配狀態,而無需額外驗證。

返回的狀態字典可能看起來像這樣

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}
返回型別

dict[str, Any]

step(closure=None)[原始碼][原始碼]

執行單個最佳化步。

引數

closure (可呼叫物件, 可選) – 一個閉包,用於重新評估模型並返回損失。

zero_grad(set_to_none=True)[原始碼]

重置所有最佳化後的 torch.Tensor 的梯度。

引數

set_to_none (布林值) – 不設定為零,而是將梯度設定為 None。這通常會降低記憶體佔用,並能適度提高效能。然而,它會改變某些行為。例如:1. 當用戶嘗試訪問梯度並對其執行手動操作時,None 屬性或填充 0 的張量會有不同的行為。2. 如果使用者請求 zero_grad(set_to_none=True) 隨後進行反向傳播,則未收到梯度的引數的 .grads 保證為 None。3. torch.optim 最佳化器在梯度為 0 或 None 時有不同的行為(一種情況下以 0 梯度執行步進,另一種情況下完全跳過步進)。

文件

查閱 PyTorch 全面的開發者文件

檢視文件

教程

獲取面向初學者和高階開發者的深入教程

檢視教程

資源

查詢開發資源並獲得問題解答

檢視資源