torch.optim.Optimizer.register_state_dict_post_hook

Optimizer.register_state_dict_post_hook(hook, prepend=False)[源代码]

注册一个状态字典的后置钩子,在调用 state_dict() 之后会被调用。

它应具有如下签名:

hook(optimizer, state_dict) -> state_dict or None

在生成selfstate_dict之后,该钩子将被调用,并使用selfstate_dict作为参数。钩子可以就地修改state_dict或选择返回一个新的state_dict。注册的钩子可以在返回之前对state_dict执行后处理。

参数
  • hook (Callable) – 用户自定义的要注册的钩子。

  • prepend (bool) – 如果为 True,则提供的 post hook 将在所有已注册的 post-hooks 之前触发。否则,提供的 hook 将在所有已注册的 post-hooks 之后触发。(默认值:False)

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemoveableHandle

本页目录