torch.optim.Optimizer.register_load_state_dict_pre_hook

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

注册一个在调用load_state_dict()之前被调用的load_state_dict预处理钩子。该钩子应具有以下签名:

hook(optimizer, state_dict) -> state_dict or None

optimizer 参数是正在使用的优化器实例,而 state_dict 参数是用户传递给 load_state_dictstate_dict 浅拷贝。钩子可以就地修改 state_dict 或者可选地返回一个新的 state_dict。如果返回了一个新的 state_dict,它将被加载到优化器中。

该钩子在调用 self.load_state_dict() 之前会被调用,并使用参数 selfstate_dict。注册的钩子可以在 load_state_dict 调用前用于执行预处理。

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

  • prepend (bool) – 如果为 True,提供的 hook 将在调用 load_state_dict 之前触发所有已注册的预处理钩子。否则,该 hook 将在所有已注册的预处理钩子之后触发。(默认值:False)

返回值

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

返回类型

torch.utils.hooks.RemoveableHandle

本页目录