deepmd.tf.model.model

Module Contents

Classes

Model	Abstract base model.
StandardModel	Standard model, which must contain a descriptor and a fitting.

class deepmd.tf.model.model.Model(type_embedding: dict | deepmd.tf.utils.type_embed.TypeEmbedNet | None = None, type_map: List[str] | None = None, data_stat_nbatch: int = 10, data_bias_nsample: int = 10, data_stat_protect: float = 0.01, use_srtab: str | None = None, smin_alpha: float | None = None, sw_rmin: float | None = None, sw_rmax: float | None = None, srtab_add_bias: bool = True, spin: deepmd.tf.utils.spin.Spin | None = None, compress: dict | None = None, **kwargs)

Bases: abc.ABC, make_plugin_registry('model')

Abstract base model.

Parameters:

type_embedding

Type embedding net

type_map

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

data_stat_nbatch

Number of frames used for data statistic

data_bias_nsample

The number of training samples in a system to compute and change the energy bias.

data_stat_protect

Protect parameter for atomic energy regression

use_srtab

The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.

smin_alpha

The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.

sw_rmin

The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

sw_rmin

The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

srtab_add_biasbool

Whether add energy bias from the statistics of the data to short-range tabulated atomic energy. It only takes effect when use_srtab is provided.

spin

spin

compress

Compression information for internal use

get_type_map() → list

Get the type map.

abstract build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = None)

Build the model.

Parameters:

coord_tf.Tensor

The coordinates of atoms

atype_tf.Tensor

The atom types of atoms

natomstf.Tensor

The number of atoms

boxtf.Tensor

The box vectors

meshtf.Tensor

The mesh vectors

input_dictdict

The input dict

frz_modelstr, optional

The path to the frozen model

ckpt_metastr, optional

The path prefix of the checkpoint and meta files

suffixstr, optional

The suffix of the scope

reusebool or tf.AUTO_REUSE, optional

Whether to reuse the variables

Returns:

dict

The output dict

init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters:

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

build_descrpt(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = None)

Build the descriptor part of the model.

Parameters:

coord_tf.Tensor

The coordinates of atoms

atype_tf.Tensor

The atom types of atoms

natomstf.Tensor

The number of atoms

boxtf.Tensor

The box vectors

meshtf.Tensor

The mesh vectors

input_dictdict

The input dict

frz_modelstr, optional

The path to the frozen model

ckpt_metastr, optional

The path prefix of the checkpoint and meta files

suffixstr, optional

The suffix of the scope

reusebool or tf.AUTO_REUSE, optional

Whether to reuse the variables

Returns:

tf.Tensor

The descriptor tensor

build_type_embedding(ntypes: int, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = None) → deepmd.tf.env.tf.Tensor

Build the type embedding part of the model.

Parameters:

ntypesint

The number of types

frz_modelstr, optional

The path to the frozen model

ckpt_metastr, optional

The path prefix of the checkpoint and meta files

suffixstr, optional

The suffix of the scope

reusebool or tf.AUTO_REUSE, optional

Whether to reuse the variables

Returns:

tf.Tensor

The type embedding tensor

_import_graph_def_from_frz_model(frz_model: str, feed_dict: dict, return_elements: List[str])

_import_graph_def_from_ckpt_meta(ckpt_meta: str, feed_dict: dict, return_elements: List[str])

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters:

mixed_precdict

The mixed precision config

change_energy_bias(data: deepmd.tf.utils.data_system.DeepmdDataSystem, frozen_model: str, origin_type_map: list, full_type_map: str, bias_adjust_mode: str = 'change-by-statistic') → None

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

Parameters:

dataDeepmdDataSystem

The training data.

frozen_modelstr

The path file of frozen model.

origin_type_maplist

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

full_type_mapstr

The full type_map in pretrained model

bias_adjust_modestr

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

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

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

enable_compression(suffix: str = '')

Enable compression.

Parameters:

suffixstr

suffix to name scope

get_numb_fparam() → int | dict

Get the number of frame parameters.

get_numb_aparam() → int | dict

Get the number of atomic parameters.

get_numb_dos() → int | dict

Get the number of gridpoints in energy space.

abstract get_fitting() → deepmd.tf.fit.fitting.Fitting | dict

Get the fitting(s).

abstract get_loss(loss: dict, lr) → deepmd.tf.loss.loss.Loss | dict | None

Get the loss function(s).

abstract get_rcut() → float

Get cutoff radius of the model.

abstract get_ntypes() → int

Get the number of types.

abstract data_stat(data: dict)

Data staticis.

get_feed_dict(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, **kwargs) → Dict[str, deepmd.tf.env.tf.Tensor]

Generate the feed_dict for current descriptor.

Parameters:

coord_tf.Tensor

The coordinate of atoms

atype_tf.Tensor

The type of atoms

natomstf.Tensor

The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

boxtf.Tensor

The box. Can be generated by deepmd.tf.model.make_stat_input

meshtf.Tensor

For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.

**kwargsdict

The additional arguments

Returns:

feed_dictdict[str, tf.Tensor]

The output feed_dict of current descriptor

abstract classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict

Update the selection and perform neighbor statistics.

Parameters:

global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

Notes

Do not modify the input data without copying it.

classmethod deserialize(data: dict, suffix: str = '') → Model

Deserialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Parameters:

datadict

The serialized data

suffixstr, optional

Name suffix to identify this model

Returns:

Model

The deserialized Model

abstract serialize(suffix: str = '') → dict

Serialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Returns:

dict

The serialized data

suffixstr, optional

Name suffix to identify this descriptor

class deepmd.tf.model.model.StandardModel(descriptor: dict | deepmd.tf.descriptor.descriptor.Descriptor, fitting_net: dict | deepmd.tf.fit.fitting.Fitting, type_embedding: dict | deepmd.tf.utils.type_embed.TypeEmbedNet | None = None, type_map: List[str] | None = None, **kwargs)

Bases: Model

Standard model, which must contain a descriptor and a fitting.

Parameters:

descriptorUnion[dict, Descriptor]

The descriptor

fitting_netUnion[dict, Fitting]

The fitting network

type_embeddingdict, optional

The type embedding

type_maplist of dict, optional

The type map

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters:

mixed_precdict

The mixed precision config

enable_compression(suffix: str = '')

Enable compression.

Parameters:

suffixstr

suffix to name scope

get_fitting() → deepmd.tf.fit.fitting.Fitting | dict

Get the fitting(s).

get_loss(loss: dict, lr) → deepmd.tf.loss.loss.Loss | dict

Get the loss function(s).

get_rcut() → float

Get cutoff radius of the model.

get_ntypes() → int

Get the number of types.

classmethod update_sel(global_jdata: dict, local_jdata: dict)

Update the selection and perform neighbor statistics.

Parameters:

global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

classmethod deserialize(data: dict, suffix: str = '') → deepmd.tf.descriptor.descriptor.Descriptor

Deserialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Parameters:

datadict

The serialized data

suffixstr, optional

Name suffix to identify this descriptor

Returns:

Descriptor

The deserialized descriptor

serialize(suffix: str = '') → dict

Serialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Returns:

dict

The serialized data

suffixstr, optional

Name suffix to identify this descriptor