deepmd.pd.train.wrapper

Attributes

_StateDict
log

Classes

ModelWrapper

Module Contents

deepmd.pd.train.wrapper._StateDict

deepmd.pd.train.wrapper.log

class deepmd.pd.train.wrapper.ModelWrapper(model: paddle.nn.Layer | dict, loss: paddle.nn.Layer | dict = None, model_params: dict[str, Any] | None = None, shared_links: dict[str, Any] | None = None)

Bases: paddle.nn.Layer

model_params

train_infos

multi_task = False

model

loss = None

inference_only

share_params(shared_links: dict[str, Any], resume: bool = False) → None

Share the parameters of classes following rules defined in shared_links during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.

forward(coord: paddle.Tensor, atype: paddle.Tensor, spin: paddle.Tensor | None = None, box: paddle.Tensor | None = None, cur_lr: paddle.Tensor | None = None, label: paddle.Tensor | None = None, task_key: paddle.Tensor | None = None, inference_only: bool = False, do_atomic_virial: bool = False, fparam: paddle.Tensor | None = None, aparam: paddle.Tensor | None = None) → dict[str, paddle.Tensor]

load_state_dict(state_dict: _StateDict) → tuple[list[str], list[str]]

set_state_dict(state_dict: _StateDict) → tuple[list[str], list[str]]

state_dict() → dict[str, Any]

set_extra_state(extra_state: dict[str, Any]) → None

get_extra_state() → dict