torchrec.optim¶
Torchrec 優化器
Torchrec 包含一個稱為 KeyedOptimizer 的特殊優化器。KeyedOptimizer 使用有意義的鍵公開 state_dict - 它允許載入 torch.tensor 和 ShardedTensor，並禁止將空狀態載入已初始化的 KeyedOptimizer，反之亦然。
它還包含：- 幾個包裝 KeyedOptimizer 的模組，稱為 CombinedOptimizer 和 OptimizerWrapper - RecSys 中使用的優化器：例如 rowwise adagrad/adam/etc
torchrec.optim.clipping¶
- class torchrec.optim.clipping.GradientClipping(value)¶
基底：
Enum列舉。
- NONE = 'none'¶
- NORM = 'norm'¶
- VALUE = 'value'¶
- class torchrec.optim.clipping.GradientClippingOptimizer(optimizer: KeyedOptimizer, clipping: GradientClipping = GradientClipping.NONE, max_gradient: float = 0.1, norm_type: Union[float, str] = 2.0)¶
基底：
OptimizerWrapper在進行優化步驟之前裁剪梯度。
- 參數：:
optimizer (KeyedOptimizer) – 要包裝的優化器
clipping (GradientClipping) – 如何裁剪梯度
max_gradient (float) – 裁剪的最大值
norm_type (float 或 str) – 使用的 p-範數的類型。對於無窮範數，可以是
'inf'。
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
torchrec.optim.fused¶
- class torchrec.optim.fused.EmptyFusedOptimizer¶
基底：
FusedOptimizer沒有操作步驟且沒有要優化的參數的融合優化器類別
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- zero_grad(set_to_none: bool = False) None¶
重設所有已優化的
torch.Tensor的梯度。- 參數：:
set_to_none (bool) – 不是設定為零，而是將梯度設定為 None。這通常會減少記憶體佔用空間，並且可以略微提高效能。但是，它會改變某些行為。例如：1. 當使用者嘗試存取梯度並對其執行手動操作時，None 屬性或充滿 0 的張量將會有不同的行為。2. 如果使用者在反向傳遞後請求
zero_grad(set_to_none=True)，則對於沒有收到梯度的參數，.grad保證為 None。3. 如果梯度為 0 或 None，torch.optim優化器會有不同的行為（在其中一種情況下，它會以 0 的梯度執行步驟，而在另一種情況下，它會完全跳過該步驟）。
- class torchrec.optim.fused.FusedOptimizer(params: Mapping[str, Union[Tensor, ShardedTensor]], state: Mapping[Any, Any], param_groups: Collection[Mapping[str, Any]])¶
基底：
KeyedOptimizer,ABC假設權重更新是在反向傳播期間完成的，因此 step() 是無效的。
- abstract step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- abstract zero_grad(set_to_none: bool = False) None¶
重設所有已優化的
torch.Tensor的梯度。- 參數：:
set_to_none (bool) – 不是設定為零，而是將梯度設定為 None。這通常會減少記憶體佔用空間，並且可以略微提高效能。但是，它會改變某些行為。例如：1. 當使用者嘗試存取梯度並對其執行手動操作時，None 屬性或充滿 0 的張量將會有不同的行為。2. 如果使用者在反向傳遞後請求
zero_grad(set_to_none=True)，則對於沒有收到梯度的參數，.grad保證為 None。3. 如果梯度為 0 或 None，torch.optim優化器會有不同的行為（在其中一種情況下，它會以 0 的梯度執行步驟，而在另一種情況下，它會完全跳過該步驟）。
- class torchrec.optim.fused.FusedOptimizerModule¶
基底：
ABC模組，在反向傳播期間進行權重更新。
- abstract property fused_optimizer: KeyedOptimizer¶
torchrec.optim.keyed¶
- class torchrec.optim.keyed.CombinedOptimizer(optims: List[Union[KeyedOptimizer, Tuple[str, KeyedOptimizer]]])¶
基底：
KeyedOptimizer將多個 KeyedOptimizers 合併為一個。
旨在為不同的子模組組合不同的優化器
- property optimizers: List[Tuple[str, KeyedOptimizer]]¶
- property param_groups: Collection[Mapping[str, Any]]¶
- property params: Mapping[str, Union[Tensor, ShardedTensor]]¶
- post_load_state_dict() None¶
- static prepend_opt_key(name: str, opt_key: str) str¶
- save_param_groups(save: bool) None¶
- property state: Mapping[Tensor, Any]¶
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- zero_grad(set_to_none: bool = False) None¶
重設所有已優化的
torch.Tensor的梯度。- 參數：:
set_to_none (bool) – 不是設定為零，而是將梯度設定為 None。這通常會減少記憶體佔用空間，並且可以略微提高效能。但是，它會改變某些行為。例如：1. 當使用者嘗試存取梯度並對其執行手動操作時，None 屬性或充滿 0 的張量將會有不同的行為。2. 如果使用者在反向傳遞後請求
zero_grad(set_to_none=True)，則對於沒有收到梯度的參數，.grad保證為 None。3. 如果梯度為 0 或 None，torch.optim優化器會有不同的行為（在其中一種情況下，它會以 0 的梯度執行步驟，而在另一種情況下，它會完全跳過該步驟）。
- class torchrec.optim.keyed.KeyedOptimizer(params: Mapping[str, Union[Tensor, ShardedTensor]], state: Mapping[Any, Any], param_groups: Collection[Mapping[str, Any]])¶
基底：
Optimizer取得參數字典，並按參數鍵公開 state_dict。
此實現比 torch.Optimizer 中的實現更嚴格：它要求實現必須在第一次優化迭代期間完全初始化其狀態，並且它禁止將空狀態加載到已初始化的 KeyedOptimizer 中，反之亦然。
預設情況下，它也不會在 state_dict() 中公開 param_groups。可以透過設定 save_param_groups 旗標來開啟舊行為。原因是在分散式訓練中，並非所有參數都存在於所有 rank 上，我們透過參數來識別 param_group。除此之外，param_groups 通常在訓練初始化期間重新設定，因此將它們儲存為狀態的一部分幾乎沒有意義。
- add_param_group(param_group: Any) None¶
將參數群組新增至
Optimizer的 param_groups。這在微調預先訓練的網路時很有用，因為凍結層可以在訓練過程中設定為可訓練並新增至
Optimizer。- 參數：:
param_group (dict) – 指定應最佳化的張量以及特定於群組的最佳化選項。
- init_state(sparse_grad_parameter_names: Optional[Set[str]] = None) None¶
執行虛擬的最佳化器步驟，允許初始化最佳化器狀態，這通常是懶惰的。這讓我們可以從檢查點就地載入最佳化器狀態。
- load_state_dict(state_dict: Mapping[str, Any]) None¶
此實現比 torch.Optimizer 中的實現更嚴格：它要求實現必須在第一次優化迭代期間完全初始化其狀態，並且它禁止將空狀態加載到已初始化的 KeyedOptimizer 中，反之亦然。
- 由於引入了嚴格性，它讓我們可以
對狀態和 param_groups 執行相容性檢查，從而提高了可用性
透過直接複製到狀態張量來避免狀態重複，例如 optimizer.step() # 確保最佳化器已初始化 sd = optimizer.state_dict() load_checkpoint(sd) # 直接將狀態複製到張量，如果需要，請重新分片 optimizer.load_state_dict(sd) # 替換 param_groups
- post_load_state_dict() None¶
- save_param_groups(save: bool) None¶
- state_dict() Dict[str, Any]¶
返回的狀態和 param_groups 將包含參數鍵，而不是 torch.Optimizer 中的參數索引。這允許實現最佳化器重新分片等進階功能。
還可以處理遵循 PyTorch 狀態協定的類和支援的資料結構。
- class torchrec.optim.keyed.KeyedOptimizerWrapper(params: Mapping[str, Union[Tensor, ShardedTensor]], optim_factory: Callable[[List[Union[Tensor, ShardedTensor]]], Optimizer])¶
基底：
KeyedOptimizer取得參數字典，並按參數鍵公開 state_dict。
方便的包裝器，用於接收 optim_factory 可呼叫對象來建立 KeyedOptimizer
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- zero_grad(set_to_none: bool = False) None¶
重設所有已優化的
torch.Tensor的梯度。- 參數：:
set_to_none (bool) – 不是設定為零，而是將梯度設定為 None。這通常會減少記憶體佔用空間，並且可以略微提高效能。但是，它會改變某些行為。例如：1. 當使用者嘗試存取梯度並對其執行手動操作時，None 屬性或充滿 0 的張量將會有不同的行為。2. 如果使用者在反向傳遞後請求
zero_grad(set_to_none=True)，則對於沒有收到梯度的參數，.grad保證為 None。3. 如果梯度為 0 或 None，torch.optim優化器會有不同的行為（在其中一種情況下，它會以 0 的梯度執行步驟，而在另一種情況下，它會完全跳過該步驟）。
- class torchrec.optim.keyed.OptimizerWrapper(optimizer: KeyedOptimizer)¶
基底：
KeyedOptimizer包裝器，接收 KeyedOptimizer 並且是一個 KeyedOptimizer
GradientClippingOptimizer 和 WarmupOptimizer 等最佳化器的子類別
- add_param_group(param_group: Any) None¶
將參數群組新增至
Optimizer的 param_groups。這在微調預先訓練的網路時很有用，因為凍結層可以在訓練過程中設定為可訓練並新增至
Optimizer。- 參數：:
param_group (dict) – 指定應最佳化的張量以及特定於群組的最佳化選項。
- load_state_dict(state_dict: Mapping[str, Any]) None¶
此實現比 torch.Optimizer 中的實現更嚴格：它要求實現必須在第一次優化迭代期間完全初始化其狀態，並且它禁止將空狀態加載到已初始化的 KeyedOptimizer 中，反之亦然。
- 由於引入了嚴格性，它讓我們可以
對狀態和 param_groups 執行相容性檢查，從而提高了可用性
透過直接複製到狀態張量來避免狀態重複，例如 optimizer.step() # 確保最佳化器已初始化 sd = optimizer.state_dict() load_checkpoint(sd) # 直接將狀態複製到張量，如果需要，請重新分片 optimizer.load_state_dict(sd) # 替換 param_groups
- post_load_state_dict() None¶
- save_param_groups(save: bool) None¶
- state_dict() Dict[str, Any]¶
返回的狀態和 param_groups 將包含參數鍵，而不是 torch.Optimizer 中的參數索引。這允許實現最佳化器重新分片等進階功能。
還可以處理遵循 PyTorch 狀態協定的類和支援的資料結構。
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- zero_grad(set_to_none: bool = False) None¶
重設所有已優化的
torch.Tensor的梯度。- 參數：:
set_to_none (bool) – 不是設定為零，而是將梯度設定為 None。這通常會減少記憶體佔用空間，並且可以略微提高效能。但是，它會改變某些行為。例如：1. 當使用者嘗試存取梯度並對其執行手動操作時，None 屬性或充滿 0 的張量將會有不同的行為。2. 如果使用者在反向傳遞後請求
zero_grad(set_to_none=True)，則對於沒有收到梯度的參數，.grad保證為 None。3. 如果梯度為 0 或 None，torch.optim優化器會有不同的行為（在其中一種情況下，它會以 0 的梯度執行步驟，而在另一種情況下，它會完全跳過該步驟）。
torchrec.optim.warmup¶
- 類別 torchrec.optim.warmup.WarmupOptimizer(optimizer: KeyedOptimizer, stages: List[WarmupStage], lr: float = 0.1, lr_param: str = 'lr', param_name: str = '__warmup')¶
基底：
OptimizerWrapper根據時間表調整學習率。
- 參數：:
optimizer (KeyedOptimizer) – 要包裝的優化器
stages (List[WarmupStage]) – 要經歷的階段
lr (float) – 初始學習率
lr_param (str) – 參數群組中的學習率參數。
param_name – 用於保存預熱狀態的虛擬參數名稱。
- post_load_state_dict() None¶
- step(closure: Optional[Any] = None) None¶
執行單一優化步驟以更新參數。
- 參數：:
closure (Callable) – 重新評估模型並返回損失的閉包。對於大多數優化器來說是可選的。
備註
除非另有說明，否則此函數不應修改參數的
.grad欄位。
- 類別 torchrec.optim.warmup.WarmupPolicy(value)¶
基底：
Enum列舉。
- CONSTANT = 'constant'¶
- COSINE_ANNEALING_WARM_RESTARTS = 'cosine_annealing_warm_restarts'¶
- INVSQRT = 'inv_sqrt'¶
- LINEAR = 'linear'¶
- NONE = 'none'¶
- POLY = 'poly'¶
- STEP = 'step'¶
- 類別 torchrec.optim.warmup.WarmupStage(policy: torchrec.optim.warmup.WarmupPolicy = <WarmupPolicy.LINEAR: 'linear'>, max_iters: int = 1, value: float = 1.0, lr_scale: float = 1.0, decay_iters: int = -1, sgdr_period: int = 1)¶
基底：
object- decay_iters: int = -1¶
- lr_scale: float = 1.0¶
- max_iters: int = 1¶
- policy: WarmupPolicy = 'linear'¶
- sgdr_period: int = 1¶
- value: float = 1.0¶
模組內容¶
Torchrec 優化器
Torchrec 包含一個稱為 KeyedOptimizer 的特殊優化器。KeyedOptimizer 使用有意義的鍵公開 state_dict - 它允許載入 torch.tensor 和 ShardedTensor，並禁止將空狀態載入已初始化的 KeyedOptimizer，反之亦然。
它還包含：- 幾個包裝 KeyedOptimizer 的模組，稱為 CombinedOptimizer 和 OptimizerWrapper - RecSys 中使用的優化器：例如 rowwise adagrad/adam/etc
