评价此页

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