評價此頁

torch.optim.Optimizer.register_state_dict_post_hook#

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

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

它應具有以下簽名

hook(optimizer, state_dict) -> state_dict or None

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

引數
  • hook (Callable) – 使用者定義的待註冊鉤子。

  • prepend (bool) – 如果為 True,則提供的後置 hook 將在 state_dict 上所有已註冊的後置鉤子之前執行。否則,提供的 hook 將在所有已註冊的後置鉤子之後執行。(預設: False)

返回

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

返回型別

torch.utils.hooks.RemoveableHandle