deepmd.pt.model.task package

deepmd.pt.model.task.BaseFitting

alias of BF

class deepmd.pt.model.task.DenoiseNet(*args, **kwargs)

Bases: DenoiseNet

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(pair_weights, diff, nlist_mask, ...)	Calculate the updated coord.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

training: bool

class deepmd.pt.model.task.DipoleFittingNet(*args, **kwargs)

Bases: GeneralFitting

Construct a dipole fitting net.

Parameters

var_namestr

The atomic property to fit, ‘dipole’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

embedding_widthint

The dimension of rotation matrix, m1.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

r_differentiable

If the variable is differentiated with respect to coordinates of atoms. Only reduciable variable are differentiable.

c_differentiable

If the variable is differentiated with respect to the cell tensor (pbc case). Only reduciable variable are differentiable.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

compute_output_stats(merged: Union[Callable[[], List[dict]], List[dict]], stat_file_path: Optional[DPPath] = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

forward(descriptor: Tensor, atype: Tensor, gr: Optional[Tensor] = None, g2: Optional[Tensor] = None, h2: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

output_def() → FittingOutputDef

Returns the output def of the fitting net.

serialize() → dict

Serialize the fitting to dict.

training: bool

class deepmd.pt.model.task.EnergyFittingNet(*args, **kwargs)

Bases: InvarFitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Based on embedding net output, alculate total energy.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

serialize() → dict

Serialize the fitting to dict.

training: bool

class deepmd.pt.model.task.EnergyFittingNetDirect(*args, **kwargs)

Bases: EnergyFittingNetDirect

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize()	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(inputs, atype[, gr, g2, h2, fparam, ...])	Based on embedding net output, alculate total energy.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

training: bool

class deepmd.pt.model.task.Fitting(*args, **kwargs)

Bases: Module, BF

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(*input)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

change_energy_bias(config, model, old_type_map: List[str], new_type_map: List[str], bias_shift='delta', ntest=10)

Change the energy bias according to the input data and the pretrained model.

Parameters

configDict

The configuration.

modelEnergyModel

Energy model loaded pre-trained model.

new_type_mapList[str]

The original type_map in dataset, they are targets to change the energy bias.

old_type_mapList[str]

The full type_map in pretrained model

bias_shiftstr

The mode for changing energy bias : [‘delta’, ‘statistic’] ‘delta’ : perform predictions on energies of target dataset,

and do least sqaure on the errors to obtain the target shift as bias.

‘statistic’ : directly use the statistic energy bias in the target dataset.

ntestint

The number of test samples in a system to change the energy bias.

share_params(base_class, shared_level, resume=False)

Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

class deepmd.pt.model.task.FittingNetAttenLcc(*args, **kwargs)

Bases: Fitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(output, pair, delta_pos, atype, ...)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(output, pair, delta_pos, atype, nframes, nloc)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class deepmd.pt.model.task.PolarFittingNet(*args, **kwargs)

Bases: GeneralFitting

Construct a polar fitting net.

Parameters

var_namestr

The atomic property to fit, ‘polar’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

embedding_widthint

The dimension of rotation matrix, m1.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

fit_diagbool

Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.

scaleList[float]

The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]

shift_diagbool

Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

compute_output_stats(merged: Union[Callable[[], List[dict]], List[dict]], stat_file_path: Optional[DPPath] = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

forward(descriptor: Tensor, atype: Tensor, gr: Optional[Tensor] = None, g2: Optional[Tensor] = None, h2: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

output_def() → FittingOutputDef

Returns the output def of the fitting net.

serialize() → dict

Serialize the fitting to dict.

training: bool

class deepmd.pt.model.task.TypePredictNet(*args, **kwargs)

Bases: Fitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(features[, masked_tokens])	Calculate the predicted logits.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(features, masked_tokens: Optional[Tensor] = None)

Calculate the predicted logits. Args: - features: Input features with shape [nframes, nloc, feature_dim]. - masked_tokens: Input masked tokens with shape [nframes, nloc].

Returns

• logits: Predicted probs with shape [nframes, nloc, ntypes].

training: bool

Submodules

deepmd.pt.model.task.atten_lcc module

class deepmd.pt.model.task.atten_lcc.FittingNetAttenLcc(*args, **kwargs)

Bases: Fitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(output, pair, delta_pos, atype, ...)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(output, pair, delta_pos, atype, nframes, nloc)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

deepmd.pt.model.task.base_fitting module

deepmd.pt.model.task.denoise module

class deepmd.pt.model.task.denoise.DenoiseNet(*args, **kwargs)

Bases: DenoiseNet

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(pair_weights, diff, nlist_mask, ...)	Calculate the updated coord.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

training: bool

deepmd.pt.model.task.dipole module

class deepmd.pt.model.task.dipole.DipoleFittingNet(*args, **kwargs)

Bases: GeneralFitting

Construct a dipole fitting net.

Parameters

var_namestr

The atomic property to fit, ‘dipole’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

embedding_widthint

The dimension of rotation matrix, m1.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

r_differentiable

If the variable is differentiated with respect to coordinates of atoms. Only reduciable variable are differentiable.

c_differentiable

If the variable is differentiated with respect to the cell tensor (pbc case). Only reduciable variable are differentiable.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

compute_output_stats(merged: Union[Callable[[], List[dict]], List[dict]], stat_file_path: Optional[DPPath] = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

forward(descriptor: Tensor, atype: Tensor, gr: Optional[Tensor] = None, g2: Optional[Tensor] = None, h2: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

output_def() → FittingOutputDef

Returns the output def of the fitting net.

serialize() → dict

Serialize the fitting to dict.

training: bool

deepmd.pt.model.task.ener module

class deepmd.pt.model.task.ener.EnergyFittingNet(*args, **kwargs)

Bases: InvarFitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Based on embedding net output, alculate total energy.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

serialize() → dict

Serialize the fitting to dict.

training: bool

class deepmd.pt.model.task.ener.EnergyFittingNetDirect(*args, **kwargs)

Bases: EnergyFittingNetDirect

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize()	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(inputs, atype[, gr, g2, h2, fparam, ...])	Based on embedding net output, alculate total energy.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

training: bool

class deepmd.pt.model.task.ener.InvarFitting(*args, **kwargs)

Bases: InvarFitting

Construct a fitting net for energy.

Parameters

var_namestr

The atomic property to fit, ‘energy’, ‘dipole’, and ‘polar’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

dim_outint

The output dimension of the fitting net.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

bias_atom_etorch.Tensor, optional

Average enery per atom for each element.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

exclude_types: List[int]

Atomic contributions of the excluded atom types are set zero.

atom_ener

Specifying atomic energy contribution in vacuum. The set_davg_zero key in the descrptor should be set.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Based on embedding net output, alculate total energy.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

training: bool

deepmd.pt.model.task.fitting module

class deepmd.pt.model.task.fitting.Fitting(*args, **kwargs)

Bases: Module, BF

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(*input)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

change_energy_bias(config, model, old_type_map: List[str], new_type_map: List[str], bias_shift='delta', ntest=10)

Change the energy bias according to the input data and the pretrained model.

Parameters

configDict

The configuration.

modelEnergyModel

Energy model loaded pre-trained model.

new_type_mapList[str]

The original type_map in dataset, they are targets to change the energy bias.

old_type_mapList[str]

The full type_map in pretrained model

bias_shiftstr

The mode for changing energy bias : [‘delta’, ‘statistic’] ‘delta’ : perform predictions on energies of target dataset,

and do least sqaure on the errors to obtain the target shift as bias.

‘statistic’ : directly use the statistic energy bias in the target dataset.

ntestint

The number of test samples in a system to change the energy bias.

share_params(base_class, shared_level, resume=False)

Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

class deepmd.pt.model.task.fitting.GeneralFitting(*args, **kwargs)

Bases: Fitting

Construct a general fitting net.

Parameters

var_namestr

The atomic property to fit, ‘energy’, ‘dipole’, and ‘polar’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

dim_outint

The output dimension of the fitting net.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

bias_atom_etorch.Tensor, optional

Average enery per atom for each element.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

exclude_types: List[int]

Atomic contributions of the excluded atom types are set zero.

trainableUnion[List[bool], bool]

If the parameters in the fitting net are trainable. Now this only supports setting all the parameters in the fitting net at one state. When in List[bool], the trainable will be True only if all the boolean parameters are True.

remove_vaccum_contribution: List[bool], optional

Remove vaccum contribution before the bias is added. The list assigned each type. For mixed_types provide [True], otherwise it should be a list of the same length as ntypes signaling if or not removing the vaccum contribution for the atom types in the list.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(*input)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

classmethod deserialize(data: dict) → GeneralFitting

Deserialize the fitting.

Parameters

datadict

The serialized data

Returns

BF

The deserialized fitting

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_sel_type() → List[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

reinit_exclude(exclude_types: List[int] = [])

serialize() → dict

Serialize the fitting to dict.

training: bool

deepmd.pt.model.task.polarizability module

class deepmd.pt.model.task.polarizability.PolarFittingNet(*args, **kwargs)

Bases: GeneralFitting

Construct a polar fitting net.

Parameters

var_namestr

The atomic property to fit, ‘polar’.

ntypesint

Element count.

dim_descrptint

Embedding width per atom.

embedding_widthint

The dimension of rotation matrix, m1.

neuronList[int]

Number of neurons in each hidden layers of the fitting net.

resnet_dtbool

Using time-step in the ResNet construction.

numb_fparamint

Number of frame parameters.

numb_aparamint

Number of atomic parameters.

activation_functionstr

Activation function.

precisionstr

Numerical precision.

mixed_typesbool

If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types.

rcondfloat, optional

The condition number for the regression of atomic energy.

seedint, optional

Random seed.

fit_diagbool

Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.

scaleList[float]

The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]

shift_diagbool

Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale.

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged[, stat_file_path])	Compute the output statistics (e.g.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(descriptor, atype[, gr, g2, h2, ...])	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_dim_aparam()	Get the number (dimension) of atomic parameters of this atomic model.
get_dim_fparam()	Get the number (dimension) of frame parameters of this atomic model.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_sel_type()	Get the selected atom types of this model.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the fitting to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
reinit_exclude

compute_output_stats(merged: Union[Callable[[], List[dict]], List[dict]], stat_file_path: Optional[DPPath] = None)

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

• List[dict]: A list of data samples from various data systems.

  Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.

• Callable[[], List[dict]]: A lazy function that returns data samples in the above format

  only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

stat_file_pathOptional[DPPath]

The path to the stat file.

forward(descriptor: Tensor, atype: Tensor, gr: Optional[Tensor] = None, g2: Optional[Tensor] = None, h2: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

output_def() → FittingOutputDef

Returns the output def of the fitting net.

serialize() → dict

Serialize the fitting to dict.

training: bool

deepmd.pt.model.task.task module

deepmd.pt.model.task.type_predict module

class deepmd.pt.model.task.type_predict.TypePredictNet(*args, **kwargs)

Bases: Fitting

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
change_energy_bias(config, model, ...[, ...])	Change the energy bias according to the input data and the pretrained model.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
compute_output_stats(merged)	Update the output bias for fitting net.
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
deserialize(data)	Deserialize the fitting.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Set the module in evaluation mode.
extra_repr()	Set the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(features[, masked_tokens])	Calculate the predicted logits.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_plugins()	Get all the registered plugins.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
init_fitting_stat(**kwargs)	Initialize the model bias by the statistics.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules()	Return an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
output_def()	Returns the output def of the fitting net.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_pre_hook(hook)	Register a pre-hook for the load_state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
serialize()	Serialize the obj to dict.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
share_memory()	See torch.Tensor.share_memory_().
share_params(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

forward(features, masked_tokens: Optional[Tensor] = None)

Calculate the predicted logits. Args: - features: Input features with shape [nframes, nloc, feature_dim]. - masked_tokens: Input masked tokens with shape [nframes, nloc].

Returns

• logits: Predicted probs with shape [nframes, nloc, ntypes].

training: bool