torch.optim.Optimizer.register_state_dict_post_hook

Optimizer.register_state_dict_post_hook(hook, prepend=False)

Register a state dict post-hook which will be called after state_dict() is called.

它应具有以下签名

hook(optimizer, state_dict) -> state_dict or None

The hook will be called with arguments self and state_dict after generating a state_dict on self. The hook may modify the state_dict inplace or optionally return a new one. The registered hook can be used to perform post-processing on the state_dict before it is returned.

参数

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

• prepend (bool) – If True, the provided post hook will be fired before all the already registered post-hooks on state_dict. Otherwise, the provided hook will be fired after all the already registered post-hooks. (default: False)

返回

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

返回类型

torch.utils.hooks.RemoveableHandle