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

钩子将在为 self 生成 state_dict 后，以参数 self 和 state_dict 调用。钩子可以原地修改 state_dict，或者选择返回一个新的 state_dict。注册的钩子可用于在 state_dict 返回之前对其进行后处理。

参数

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

• prepend (bool) – 如果为 True，提供的后置 hook 将在 state_dict 上所有已注册的后置钩子之前触发。否则，提供的 hook 将在所有已注册的后置钩子之后触发。(默认为: False)

返回

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

返回类型

torch.utils.hooks.RemoveableHandle