deepmd.tf.utils.serialization
Attributes
Exceptions
Common base class for all non-exit exceptions. |
Classes
Abstract base model. |
Functions
| Load tensor value from the given tf.Graph object. |
| Load graph as well as the graph_def from the frozen model(model_file). |
| Run session with erorrs caught. |
| Serialize the model file to a dictionary. |
| Deserialize the dictionary to a model file. |
Module Contents
- class deepmd.tf.utils.serialization.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, spin: deepmd.tf.utils.spin.Spin | None = None, compress: dict | None = None, **kwargs)[source]
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
- 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)[source]
Build the model.
- Parameters:
- coord_
tf.Tensor
The coordinates of atoms
- atype_
tf.Tensor
The atom types of atoms
- natoms
tf.Tensor
The number of atoms
- box
tf.Tensor
The box vectors
- mesh
tf.Tensor
The mesh vectors
- input_dict
dict
The input dict
- frz_model
str
,optional
The path to the frozen model
- ckpt_meta
str
,optional
The path prefix of the checkpoint and meta files
- suffix
str
,optional
The suffix of the scope
- reusebool or
tf.AUTO_REUSE
,optional
Whether to reuse the variables
- coord_
- 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 [source]
Init the embedding net variables with the given frozen model.
- 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)[source]
Build the descriptor part of the model.
- Parameters:
- coord_
tf.Tensor
The coordinates of atoms
- atype_
tf.Tensor
The atom types of atoms
- natoms
tf.Tensor
The number of atoms
- box
tf.Tensor
The box vectors
- mesh
tf.Tensor
The mesh vectors
- input_dict
dict
The input dict
- frz_model
str
,optional
The path to the frozen model
- ckpt_meta
str
,optional
The path prefix of the checkpoint and meta files
- suffix
str
,optional
The suffix of the scope
- reusebool or
tf.AUTO_REUSE
,optional
Whether to reuse the variables
- coord_
- 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 [source]
Build the type embedding part of the model.
- Parameters:
- Returns:
tf.Tensor
The type embedding tensor
- _import_graph_def_from_frz_model(frz_model: str, feed_dict: dict, return_elements: List[str])[source]
- _import_graph_def_from_ckpt_meta(ckpt_meta: str, feed_dict: dict, return_elements: List[str])[source]
- enable_mixed_precision(mixed_prec: dict)[source]
Enable mixed precision for the model.
- Parameters:
- mixed_prec
dict
The mixed precision config
- mixed_prec
- 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 [source]
Change the energy bias according to the input data and the pretrained model.
- Parameters:
- data
DeepmdDataSystem
The training data.
- frozen_model
str
The path file of frozen model.
- origin_type_map
list
The original type_map in dataset, they are targets to change the energy bias.
- full_type_map
str
The full type_map in pretrained model
- bias_adjust_mode
str
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.
- data
- enable_compression(suffix: str = '')[source]
Enable compression.
- Parameters:
- suffix
str
suffix to name scope
- suffix
- abstract get_fitting() deepmd.tf.fit.fitting.Fitting | dict [source]
Get the fitting(s).
- abstract get_loss(loss: dict, lr) deepmd.tf.loss.loss.Loss | dict | None [source]
Get the loss function(s).
- 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] [source]
Generate the feed_dict for current descriptor.
- Parameters:
- coord_
tf.Tensor
The coordinate of atoms
- atype_
tf.Tensor
The type of atoms
- natoms
tf.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
- box
tf.Tensor
The box. Can be generated by deepmd.tf.model.make_stat_input
- mesh
tf.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.
- **kwargs
dict
The additional arguments
- coord_
- Returns:
- classmethod update_sel(train_data: deepmd.tf.utils.data_system.DeepmdDataSystem, type_map: List[str] | None, local_jdata: dict) Tuple[dict, float | None] [source]
- Abstractmethod:
Update the selection and perform neighbor statistics.
- Parameters:
- Returns:
Notes
Do not modify the input data without copying it.
- classmethod deserialize(data: dict, suffix: str = '') Model [source]
Deserialize the model.
There is no suffix in a native DP model, but it is important for the TF backend.
- abstract serialize(suffix: str = '') dict [source]
Serialize the model.
There is no suffix in a native DP model, but it is important for the TF backend.
- property input_requirement: List[deepmd.utils.data.DataRequirementItem]
- Abstractmethod:
Return data requirements needed for the model input.
- exception deepmd.tf.utils.serialization.GraphWithoutTensorError[source]
Bases:
Exception
Common base class for all non-exit exceptions.
- deepmd.tf.utils.serialization.get_tensor_by_name_from_graph(graph: deepmd.tf.env.tf.Graph, tensor_name: str) deepmd.tf.env.tf.Tensor [source]
Load tensor value from the given tf.Graph object.
- Parameters:
- Returns:
tf.Tensor
The tensor which was loaded from the frozen model
- Raises:
GraphWithoutTensorError
Whether the tensor_name is within the frozen model
- deepmd.tf.utils.serialization.load_graph_def(model_file: str) Tuple[deepmd.tf.env.tf.Graph, deepmd.tf.env.tf.GraphDef] [source]
Load graph as well as the graph_def from the frozen model(model_file).
- deepmd.tf.utils.serialization.run_sess(sess: deepmd.tf.env.tf.Session, *args, **kwargs)[source]
Run session with erorrs caught.
- Parameters:
- sess
tf.Session
TensorFlow Session
- *args
Variable length argument list.
- **kwargs
Arbitrary keyword arguments.
- sess
- Returns:
Any
the result of sess.run()