deepmd.pt.model.backbone.evoformer2b

Module Contents

Classes

Evoformer2bBackBone

Base class for all neural network modules.
class deepmd.pt.model.backbone.evoformer2b.Evoformer2bBackBone(nnei, layer_num=6, attn_head=8, atomic_dim=1024, pair_dim=100, feature_dim=1024, ffn_dim=2048, post_ln=False, final_layer_norm=True, final_head_layer_norm=False, emb_layer_norm=False, atomic_residual=False, evo_residual=False, residual_factor=1.0, activation_function='gelu', **kwargs)[source]

Bases: deepmd.pt.model.backbone.BackBone

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.

forward(atomic_rep, pair_rep, nlist, nlist_type, nlist_mask)[source]

Encoder the atomic and pair representations.

Args: - atomic_rep: Atomic representation with shape [nframes, nloc, atomic_dim]. - pair_rep: Pair representation with shape [nframes, nloc, nnei, pair_dim]. - nlist: Neighbor list with shape [nframes, nloc, nnei]. - nlist_type: Neighbor types with shape [nframes, nloc, nnei]. - nlist_mask: Neighbor mask with shape [nframes, nloc, nnei], False if blank.

Returns:
  • atomic_rep: Atomic representation after encoder with shape [nframes, nloc, feature_dim].
  • transformed_atomic_rep: Transformed atomic representation after encoder with shape [nframes, nloc, atomic_dim].
  • pair_rep: Pair representation after encoder with shape [nframes, nloc, nnei, attn_head].
  • delta_pair_rep: Delta pair representation after encoder with shape [nframes, nloc, nnei, attn_head].
  • norm_x: Normalization loss of atomic_rep.
  • norm_delta_pair_rep: Normalization loss of delta_pair_rep.