Python PyTorch DistributedModelParallel.state_dict用法及代碼示例

本文簡要介紹python語言中 torchrec.distributed.model_parallel.DistributedModelParallel.state_dict 的用法。

用法:

state_dict(destination: Optional[Dict[str, Any]] = None, prefix: str = '', keep_vars: bool = False) → Dict[str, Any]

返回：

包含模塊整個狀態的字典

返回類型：

dict

返回包含模塊整個狀態的字典。

包括參數和持久緩衝區(例如運行平均值)。鍵是對應的參數和緩衝區名稱。不包括設置為None 的參數和緩衝區。

例子：

>>> module.state_dict().keys()
['bias', 'weight']