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: Any, kalman_lambda: float = 0.98, kalman_nue: float = 0.9987, block_size: int = 5120)

Bases: torch.optim.optimizer.Optimizer

_params

_state

dist_init

rank

dindex = []

remainder = 0

__init_P() → None

__get_blocksize() → int

__get_nue() → float

__split_weights(weight: torch.Tensor) → list[torch.Tensor]

__update(H: torch.Tensor, error: torch.Tensor, weights: torch.Tensor) → None

set_grad_prefactor(grad_prefactor: float) → None

step(error: torch.Tensor) → None

get_device_id(index: int) → int | None