core.models.equiformer_v2.heads

core.models.equiformer_v2.heads#

Submodules#

Classes#

`Rank2SymmetricTensorHead`	A rank 2 symmetric tensor prediction head.
`EqV2ScalarHead`	Base class for all neural network modules.
`EqV2VectorHead`	Base class for all neural network modules.

Package Contents#

class core.models.equiformer_v2.heads.Rank2SymmetricTensorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'stress', decompose: bool = False, edge_level_mlp: bool = False, num_mlp_layers: int = 2, use_source_target_embedding: bool = False, extensive: bool = False, avg_num_nodes: int = 1.0, default_norm_type: str = 'layer_norm_sh')#

Bases: torch.nn.Module, fairchem.core.models.base.HeadInterface

A rank 2 symmetric tensor prediction head.

ouput_name#: name of output prediction property (ie, stress)

sphharm_norm#: layer normalization for spherical harmonic edge weights

xedge_layer_norm#: embedding layer norm

block#: rank 2 equivariant symmetric tensor block

output_name#

decompose#

use_source_target_embedding#

avg_num_nodes#

sphharm_norm#

xedge_layer_norm#

forward(data: dict[str, torch.Tensor] | torch.Tensor, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Parameters:

data – data batch
emb – dictionary with embedding object and graph data

Returns: dict of {output property name: predicted value}

class core.models.equiformer_v2.heads.EqV2ScalarHead(backbone, output_name: str = 'energy', reduce: str = 'sum')#

Bases: torch.nn.Module, fairchem.core.models.base.HeadInterface

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.

output_name#

reduce#

avg_num_nodes#

energy_block#

forward(data: torch_geometric.data.Batch, emb: dict[str, torch.Tensor | fairchem.core.models.base.GraphData])#

Head forward.

Parameters:

data (DataBatch) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

class core.models.equiformer_v2.heads.EqV2VectorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'forces')#

Bases: torch.nn.Module, fairchem.core.models.base.HeadInterface

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.

output_name#

activation_checkpoint#

force_block#

forward(data: torch_geometric.data.Batch, emb: dict[str, torch.Tensor])#

Head forward.

Parameters:

data (DataBatch) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]