core.models.base#
Copyright (c) Meta Platforms, Inc. and affiliates.
This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.
Classes#
Base class for all neural network modules. |
|
Base class for all neural network modules. |
Module Contents#
- class core.models.base.HeadInterface#
- property use_amp#
- abstract forward(data: fairchem.core.datasets.atomic_data.AtomicData, emb: dict[str, torch.Tensor]) dict[str, torch.Tensor] #
Head forward.
- Parameters:
data (AtomicData) – 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.base.BackboneInterface#
- abstract forward(data: fairchem.core.datasets.atomic_data.AtomicData) dict[str, torch.Tensor] #
Backbone forward.
- Parameters:
data (AtomicData) – Atomic systems as input
- Returns:
embedding – Return backbone embeddings for the given input
- Return type:
dict[str->torch.Tensor]
- class core.models.base.HydraModel(backbone: dict | None = None, heads: dict | None = None, finetune_config: dict | None = None, otf_graph: bool = True, pass_through_head_outputs: bool = False, freeze_backbone: bool = False)#
Bases:
torch.nn.Module
Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested 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) -> None: 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 also have their parameters converted 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.
- device = None#
- otf_graph#
- pass_through_head_outputs#
- forward(data: fairchem.core.datasets.atomic_data.AtomicData)#
- class core.models.base.HydraModelV2(backbone: BackboneInterface, heads: dict[str, HeadInterface], freeze_backbone: bool = False)#
Bases:
torch.nn.Module
Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested 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) -> None: 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 also have their parameters converted 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.
- backbone#
- output_heads#
- device = None#
- forward(data)#