deepmd.pd.train.wrapper

deepmd.pd.train.wrapper#

Attributes#

`_StateDict`
`log`

Classes#

ModelWrapper

Module Contents#

deepmd.pd.train.wrapper._StateDict[source]#

deepmd.pd.train.wrapper.log[source]#

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

Bases: paddle.nn.Layer

model_params = None[source]#

train_infos[source]#

multi_task = False[source]#

model[source]#

loss = None[source]#

inference_only[source]#

share_params(shared_links, resume=False) → None[source]#: 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, atype, 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=False, do_atomic_virial=False, fparam: paddle.Tensor | None = None, aparam: paddle.Tensor | None = None)[source]#

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

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

state_dict()[source]#

set_extra_state(extra_state: dict)[source]#

get_extra_state() → dict[source]#