deepmd.pt.optimizer

Submodules

• deepmd.pt.optimizer.KFWrapper
• deepmd.pt.optimizer.LKF

Classes

KFOptimizerWrapper
LKFOptimizer

Package Contents

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

model

optimizer

atoms_selected

atoms_per_group

is_distributed = False

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

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

update_denoise_coord(inputs: dict, clean_coord: torch.Tensor, update_prefactor: float = 1, mask_loss_coord: bool = True, coord_mask: torch.Tensor = None) → None

__sample(atoms_selected: int, atoms_per_group: int, natoms: int) → numpy.ndarray

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

Bases: torch.optim.optimizer.Optimizer

_params

_state

dist_init

rank

dindex = []

remainder = 0

__init_P() → None

__get_blocksize()

__get_nue()

__split_weights(weight)

__update(H, error, weights) → None

set_grad_prefactor(grad_prefactor) → None

step(error) → None

get_device_id(index)