deepmd.pt.optimizer package

class deepmd.pt.optimizer.KFOptimizerWrapper(model: Module, optimizer: Optimizer, atoms_selected: int, atoms_per_group: int, is_distributed: bool = False)

Bases: object

Methods

update_denoise_coord
update_energy
update_force

update_denoise_coord(inputs: dict, clean_coord: Tensor, update_prefactor: float = 1, mask_loss_coord: bool = True, coord_mask: Optional[Tensor] = None) → None

update_energy(inputs: dict, Etot_label: Tensor, update_prefactor: float = 1) → None

update_force(inputs: dict, Force_label: Tensor, update_prefactor: float = 1) → None

class deepmd.pt.optimizer.LKFOptimizer(params, kalman_lambda=0.98, kalman_nue=0.9987, block_size=5120)

Bases: Optimizer

Methods

add_param_group(param_group)	Add a param group to the Optimizer s param_groups.
load_state_dict(state_dict)	Loads the optimizer state.
register_load_state_dict_post_hook(hook[, ...])	Register a load_state_dict post-hook which will be called after load_state_dict() is called. It should have the following signature::.
register_load_state_dict_pre_hook(hook[, ...])	Register a load_state_dict pre-hook which will be called before load_state_dict() is called. It should have the following signature::.
register_state_dict_post_hook(hook[, prepend])	Register a state dict post-hook which will be called after state_dict() is called. It should have the following signature::.
register_state_dict_pre_hook(hook[, prepend])	Register a state dict pre-hook which will be called before state_dict() is called. It should have the following signature::.
register_step_post_hook(hook)	Register an optimizer step post hook which will be called after optimizer step. It should have the following signature::.
register_step_pre_hook(hook)	Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature::.
state_dict()	Returns the state of the optimizer as a dict.
step(error)	Performs a single optimization step (parameter update).
zero_grad([set_to_none])	Resets the gradients of all optimized torch.Tensor s.

OptimizerPostHook
OptimizerPreHook
profile_hook_step
set_grad_prefactor

set_grad_prefactor(grad_prefactor)

step(error)

Performs a single optimization step (parameter update).

Parameters

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

Submodules

deepmd.pt.optimizer.KFWrapper module

class deepmd.pt.optimizer.KFWrapper.KFOptimizerWrapper(model: Module, optimizer: Optimizer, atoms_selected: int, atoms_per_group: int, is_distributed: bool = False)

Bases: object

Methods

update_denoise_coord
update_energy
update_force

update_denoise_coord(inputs: dict, clean_coord: Tensor, update_prefactor: float = 1, mask_loss_coord: bool = True, coord_mask: Optional[Tensor] = None) → None

update_energy(inputs: dict, Etot_label: Tensor, update_prefactor: float = 1) → None

update_force(inputs: dict, Force_label: Tensor, update_prefactor: float = 1) → None

deepmd.pt.optimizer.LKF module

class deepmd.pt.optimizer.LKF.LKFOptimizer(params, kalman_lambda=0.98, kalman_nue=0.9987, block_size=5120)

Bases: Optimizer

Methods

add_param_group(param_group)	Add a param group to the Optimizer s param_groups.
load_state_dict(state_dict)	Loads the optimizer state.
register_load_state_dict_post_hook(hook[, ...])	Register a load_state_dict post-hook which will be called after load_state_dict() is called. It should have the following signature::.
register_load_state_dict_pre_hook(hook[, ...])	Register a load_state_dict pre-hook which will be called before load_state_dict() is called. It should have the following signature::.
register_state_dict_post_hook(hook[, prepend])	Register a state dict post-hook which will be called after state_dict() is called. It should have the following signature::.
register_state_dict_pre_hook(hook[, prepend])	Register a state dict pre-hook which will be called before state_dict() is called. It should have the following signature::.
register_step_post_hook(hook)	Register an optimizer step post hook which will be called after optimizer step. It should have the following signature::.
register_step_pre_hook(hook)	Register an optimizer step pre hook which will be called before optimizer step. It should have the following signature::.
state_dict()	Returns the state of the optimizer as a dict.
step(error)	Performs a single optimization step (parameter update).
zero_grad([set_to_none])	Resets the gradients of all optimized torch.Tensor s.

OptimizerPostHook
OptimizerPreHook
profile_hook_step
set_grad_prefactor

set_grad_prefactor(grad_prefactor)

step(error)

Performs a single optimization step (parameter update).

Parameters

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.