deepmd.pt.model.descriptor.dpa2
Module Contents
Classes
Base class for all neural network modules. |
- class deepmd.pt.model.descriptor.dpa2.DescrptDPA2(ntypes: int, repinit_rcut: float, repinit_rcut_smth: float, repinit_nsel: int, repformer_rcut: float, repformer_rcut_smth: float, repformer_nsel: int, tebd_dim: int = 8, concat_output_tebd: bool = True, repinit_neuron: List[int] = [25, 50, 100], repinit_axis_neuron: int = 16, repinit_set_davg_zero: bool = True, repinit_activation='tanh', repformer_nlayers: int = 3, repformer_g1_dim: int = 128, repformer_g2_dim: int = 16, repformer_axis_dim: int = 4, repformer_do_bn_mode: str = 'no', repformer_bn_momentum: float = 0.1, repformer_update_g1_has_conv: bool = True, repformer_update_g1_has_drrd: bool = True, repformer_update_g1_has_grrg: bool = True, repformer_update_g1_has_attn: bool = True, repformer_update_g2_has_g1g1: bool = True, repformer_update_g2_has_attn: bool = True, repformer_update_h2: bool = False, repformer_attn1_hidden: int = 64, repformer_attn1_nhead: int = 4, repformer_attn2_hidden: int = 16, repformer_attn2_nhead: int = 4, repformer_attn2_has_gate: bool = False, repformer_activation: str = 'tanh', repformer_update_style: str = 'res_avg', repformer_set_davg_zero: bool = True, repformer_add_type_ebd_to_seq: bool = False, env_protection: float = 0.0, trainable: bool = True, exclude_types: List[Tuple[int, int]] = [], type: str | None = None, rcut: float | None = None, rcut_smth: float | None = None, sel: int | None = None)[source]
Bases:
torch.nn.Module
,deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor
Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:
import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))
Submodules assigned in this way will be registered, and will have their parameters converted too when you call
to()
, etc.Note
As per the example above, an
__init__()
call to the parent class must be made before assignment on the child.- Variables:
training (bool) – Boolean represents whether this module is in training or evaluation mode.
- mixed_types() bool [source]
If true, the discriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.
If false, the discriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.
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.
- compute_input_stats(merged: Callable[[], List[dict]] | List[dict], path: deepmd.utils.path.DPPath | None = None)[source]
Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
- Parameters:
- merged
Union
[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.
- path
Optional
[DPPath
] The path to the stat file.
- merged
- abstract classmethod deserialize() DescrptDPA2 [source]
Deserialize from a dict.
- forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None)[source]
Compute the descriptor.
- Parameters:
- coord_ext
The extended coordinates of atoms. shape: nf x (nallx3)
- atype_ext
The extended aotm types. shape: nf x nall
- nlist
The neighbor list. shape: nf x nloc x nnei
- mapping
The index mapping, mapps extended region index to local region.
- comm_dict
The data needed for communication for parallel inference.
- Returns:
descriptor
The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr
The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2
The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2
The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw
The smooth switch function. shape: nf x nloc x nnei