Move select models to backbone + heads format and add support for hydra

src/fairchem/core/models/dimenet_plus_plus.py

@@ -241,7 +241,6 @@ def __init__(
         act = activation_resolver(act)
 
         super().__init__()
-
         self.cutoff = cutoff
 
         if sym is None:
@@ -330,6 +329,33 @@ def forward(self, z, pos, batch=None):
         raise NotImplementedError
 
 
+@registry.register_model("dimenetplusplus_energy_and_force_head")
+class DimeNetPlusPlusWrap_energy_and_force_head(nn.Module):
+    def __init__(self, backbone, backbone_config, head_config):
+        super().__init__()
+        self.regress_forces = backbone.regress_forces
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, x, emb):
+        outputs = {
+            "energy": (
+                emb["P"].sum(dim=0)
+                if x.batch is None
+                else scatter(emb["P"], x.batch, dim=0)
+            )
+        }
+        if self.regress_forces:
+            outputs["forces"] = -1 * (
+                torch.autograd.grad(
+                    outputs["energy"],
+                    x.pos,
+                    grad_outputs=torch.ones_like(outputs["energy"]),
+                    create_graph=True,
+                )[0]
+            )
+        return outputs
+
+
 @registry.register_model("dimenetplusplus")
 class DimeNetPlusPlusWrap(DimeNetPlusPlus, BaseModel):
     def __init__(
@@ -441,16 +467,13 @@ def forward(self, data):
         outputs = {"energy": energy}
 
         if self.regress_forces:
-            forces = (
-                -1
-                * (
-                    torch.autograd.grad(
-                        energy,
-                        data.pos,
-                        grad_outputs=torch.ones_like(energy),
-                        create_graph=True,
-                    )[0]
-                )
+            forces = -1 * (
+                torch.autograd.grad(
+                    energy,
+                    data.pos,
+                    grad_outputs=torch.ones_like(energy),
+                    create_graph=True,
+                )[0]
             )
             outputs["forces"] = forces
 
@@ -459,3 +482,66 @@ def forward(self, data):
     @property
     def num_params(self) -> int:
         return sum(p.numel() for p in self.parameters())
+
+
+@registry.register_model("dimenetplusplus_backbone")
+class DimeNetPlusPlusWrapBB(DimeNetPlusPlusWrap):
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data):
+        if self.regress_forces:
+            data.pos.requires_grad_(True)
+        pos = data.pos
+        (
+            edge_index,
+            dist,
+            _,
+            cell_offsets,
+            offsets,
+            neighbors,
+        ) = self.generate_graph(data)
+
+        data.edge_index = edge_index
+        data.cell_offsets = cell_offsets
+        data.neighbors = neighbors
+        j, i = edge_index
+
+        _, _, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets(
+            edge_index,
+            data.cell_offsets,
+            num_nodes=data.atomic_numbers.size(0),
+        )
+
+        # Calculate angles.
+        pos_i = pos[idx_i].detach()
+        pos_j = pos[idx_j].detach()
+        if self.use_pbc:
+            pos_ji, pos_kj = (
+                pos[idx_j].detach() - pos_i + offsets[idx_ji],
+                pos[idx_k].detach() - pos_j + offsets[idx_kj],
+            )
+        else:
+            pos_ji, pos_kj = (
+                pos[idx_j].detach() - pos_i,
+                pos[idx_k].detach() - pos_j,
+            )
+
+        a = (pos_ji * pos_kj).sum(dim=-1)
+        b = torch.cross(pos_ji, pos_kj).norm(dim=-1)
+        angle = torch.atan2(b, a)
+
+        rbf = self.rbf(dist)
+        sbf = self.sbf(dist, angle, idx_kj)
+
+        # Embedding block.
+        x = self.emb(data.atomic_numbers.long(), rbf, i, j)
+        P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0))
+
+        # Interaction blocks.
+        for interaction_block, output_block in zip(
+            self.interaction_blocks, self.output_blocks[1:]
+        ):
+            x = interaction_block(x, rbf, sbf, idx_kj, idx_ji)
+            P += output_block(x, rbf, i, num_nodes=pos.size(0))
+
+        return {"P": P, "edge_embedding": x, "edge_idx": i}

src/fairchem/core/models/equiformer_v2/equiformer_v2_oc20.py

@@ -678,3 +678,225 @@ def no_weight_decay(self) -> set:
                     no_wd_list.append(global_parameter_name)
 
         return set(no_wd_list)
+
+
+@registry.register_model("equiformer_v2_backbone")
+class EquiformerV2_OC20BB(EquiformerV2_OC20):
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data):
+        self.batch_size = len(data.natoms)
+        self.dtype = data.pos.dtype
+        self.device = data.pos.device
+        atomic_numbers = data.atomic_numbers.long()
+
+        (
+            edge_index,
+            edge_distance,
+            edge_distance_vec,
+            cell_offsets,
+            _,  # cell offset distances
+            neighbors,
+        ) = self.generate_graph(
+            data,
+            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+        )
+
+        data_batch_full = data.batch
+        data_batch = data.batch
+        atomic_numbers_full = atomic_numbers
+        node_offset = 0
+        if gp_utils.initialized():
+            (
+                atomic_numbers,
+                data_batch,
+                node_offset,
+                edge_index,
+                edge_distance,
+                edge_distance_vec,
+            ) = self._init_gp_partitions(
+                atomic_numbers_full,
+                data_batch_full,
+                edge_index,
+                edge_distance,
+                edge_distance_vec,
+            )
+        ###############################################################
+        # Entering Graph Parallel Region
+        # after this point, if using gp, then node, edge tensors are split
+        # across the graph parallel ranks, some full tensors such as
+        # atomic_numbers_full are required because we need to index into the
+        # full graph when computing edge embeddings or reducing nodes from neighbors
+        #
+        # all tensors that do not have the suffix "_full" refer to the partial tensors.
+        # if not using gp, the full values are equal to the partial values
+        # ie: atomic_numbers_full == atomic_numbers
+        ###############################################################
+
+        ###############################################################
+        # Initialize data structures
+        ###############################################################
+
+        # Compute 3x3 rotation matrix per edge
+        edge_rot_mat = self._init_edge_rot_mat(data, edge_index, edge_distance_vec)
+
+        # Initialize the WignerD matrices and other values for spherical harmonic calculations
+        for i in range(self.num_resolutions):
+            self.SO3_rotation[i].set_wigner(edge_rot_mat)
+
+        ###############################################################
+        # Initialize node embeddings
+        ###############################################################
+
+        # Init per node representations using an atomic number based embedding
+        offset = 0
+        x = SO3_Embedding(
+            len(atomic_numbers),
+            self.lmax_list,
+            self.sphere_channels,
+            self.device,
+            self.dtype,
+        )
+
+        offset_res = 0
+        offset = 0
+        # Initialize the l = 0, m = 0 coefficients for each resolution
+        for i in range(self.num_resolutions):
+            if self.num_resolutions == 1:
+                x.embedding[:, offset_res, :] = self.sphere_embedding(atomic_numbers)
+            else:
+                x.embedding[:, offset_res, :] = self.sphere_embedding(atomic_numbers)[
+                    :, offset : offset + self.sphere_channels
+                ]
+            offset = offset + self.sphere_channels
+            offset_res = offset_res + int((self.lmax_list[i] + 1) ** 2)
+
+        # Edge encoding (distance and atom edge)
+        edge_distance = self.distance_expansion(edge_distance)
+        if self.share_atom_edge_embedding and self.use_atom_edge_embedding:
+            source_element = atomic_numbers_full[
+                edge_index[0]
+            ]  # Source atom atomic number
+            target_element = atomic_numbers_full[
+                edge_index[1]
+            ]  # Target atom atomic number
+            source_embedding = self.source_embedding(source_element)
+            target_embedding = self.target_embedding(target_element)
+            edge_distance = torch.cat(
+                (edge_distance, source_embedding, target_embedding), dim=1
+            )
+
+        # Edge-degree embedding
+        edge_degree = self.edge_degree_embedding(
+            atomic_numbers_full,
+            edge_distance,
+            edge_index,
+            len(atomic_numbers),
+            node_offset,
+        )
+        x.embedding = x.embedding + edge_degree.embedding
+
+        ###############################################################
+        # Update spherical node embeddings
+        ###############################################################
+
+        for i in range(self.num_layers):
+            x = self.blocks[i](
+                x,  # SO3_Embedding
+                atomic_numbers_full,
+                edge_distance,
+                edge_index,
+                batch=data_batch,  # for GraphDropPath
+                node_offset=node_offset,
+            )
+
+        # Final layer norm
+        x.embedding = self.norm(x.embedding)
+
+        return {
+            "node_embedding": x,
+            "edge_distance": edge_distance,
+            "edge_index": edge_index,
+            # returning this only because it's cast to long and
+            # we don't want to repeat this.
+            "atomic_numbers": atomic_numbers_full,
+            # TODO: this is only used by graph parallel to split up the partitions,
+            # should figure out cleaner way to pass this around to the heads
+            "node_offset": node_offset,
+        }
+
+
+@registry.register_model("equiformer_v2_energy_head")
+class EquiformerV2_OC20_energy_head(nn.Module):
+    def __init__(self, backbone, backbone_config, head_config):
+        super().__init__()
+        self.avg_num_nodes = backbone.avg_num_nodes
+        self.energy_block = FeedForwardNetwork(
+            backbone.sphere_channels,
+            backbone.ffn_hidden_channels,
+            1,
+            backbone.lmax_list,
+            backbone.mmax_list,
+            backbone.SO3_grid,
+            backbone.ffn_activation,
+            backbone.use_gate_act,
+            backbone.use_grid_mlp,
+            backbone.use_sep_s2_act,
+        )
+
+    def forward(self, x, emb):
+        node_energy = self.energy_block(emb["node_embedding"])
+        node_energy = node_energy.embedding.narrow(1, 0, 1)
+        if gp_utils.initialized():
+            node_energy = gp_utils.gather_from_model_parallel_region(node_energy, dim=0)
+        energy = torch.zeros(
+            len(x.natoms),
+            device=node_energy.device,
+            dtype=node_energy.dtype,
+        )
+        energy.index_add_(0, x.batch, node_energy.view(-1))
+        return {"energy": energy / self.avg_num_nodes}
+
+
+@registry.register_model("equiformer_v2_force_head")
+class EquiformerV2_OC20_force_head(nn.Module):
+    def __init__(self, backbone, backbone_config, head_config):
+        super().__init__()
+
+        self.force_block = SO2EquivariantGraphAttention(
+            backbone.sphere_channels,
+            backbone.attn_hidden_channels,
+            backbone.num_heads,
+            backbone.attn_alpha_channels,
+            backbone.attn_value_channels,
+            1,
+            backbone.lmax_list,
+            backbone.mmax_list,
+            backbone.SO3_rotation,
+            backbone.mappingReduced,
+            backbone.SO3_grid,
+            backbone.max_num_elements,
+            backbone.edge_channels_list,
+            backbone.block_use_atom_edge_embedding,
+            backbone.use_m_share_rad,
+            backbone.attn_activation,
+            backbone.use_s2_act_attn,
+            backbone.use_attn_renorm,
+            backbone.use_gate_act,
+            backbone.use_sep_s2_act,
+            alpha_drop=0.0,
+        )
+
+    def forward(self, x, emb):
+        forces = self.force_block(
+            emb["node_embedding"],
+            x.atomic_numbers.long(),
+            emb["edge_distance"],
+            emb["edge_index"],
+            node_offset=emb["node_offset"],
+        )
+        forces = forces.embedding.narrow(1, 1, 3)
+        forces = forces.view(-1, 3).contiguous()
+        if gp_utils.initialized():
+            forces = gp_utils.gather_from_model_parallel_region(forces, dim=0)
+        return {"forces": forces}