deepmd.tf.utils.type_embed
Module Contents
Classes
Type embedding network. |
Functions
| Make the embedded type for the atoms in system. |
- deepmd.tf.utils.type_embed.embed_atom_type(ntypes: int, natoms: deepmd.tf.env.tf.Tensor, type_embedding: deepmd.tf.env.tf.Tensor)[source]
Make the embedded type for the atoms in system. The atoms are assumed to be sorted according to the type, thus their types are described by a tf.Tensor natoms, see explanation below.
- Parameters:
- ntypes:
Number of types.
- natoms:
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
- type_embedding:
The type embedding. It has the shape of [ntypes, embedding_dim]
- Returns:
atom_embeddingThe embedded type of each atom. It has the shape of [numb_atoms, embedding_dim]
- class deepmd.tf.utils.type_embed.TypeEmbedNet(*, ntypes: int, neuron: List[int], resnet_dt: bool = False, activation_function: str | None = 'tanh', precision: str = 'default', trainable: bool = True, seed: int | None = None, uniform_seed: bool = False, padding: bool = False, **kwargs)[source]
Type embedding network.
- Parameters:
- ntypes
int Number of atom types
- neuron
list[int] Number of neurons in each hidden layers of the embedding net
- resnet_dt
Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
- activation_function
The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.
- precision
The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.
- trainable
If the weights of embedding net are trainable.
- seed
Random seed for initializing the network parameters.
- uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
- padding
Concat the zero padding to the output, as the default embedding of empty type.
- ntypes
- build(ntypes: int, reuse=None, suffix='')[source]
Build the computational graph for the descriptor.
- Parameters:
- ntypes
Number of atom types.
- reuse
The weights in the networks should be reused when get the variable.
- suffix
Name suffix to identify this descriptor
- Returns:
embedded_typesThe computational graph for embedded types
- init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix='', model_type='original_model') None[source]
Init the type embedding net variables with the given dict.
- Parameters:
- graph
tf.Graph The input frozen model graph
- graph_def
tf.GraphDef The input frozen model graph_def
- suffix
Name suffix to identify this descriptor
- model_type
Indicator of whether this model is a compressed model
- graph