spamosaic.architectures.hgt

class spamosaic.architectures.hgt.HGT(*args: Any, **kwargs: Any)[source]

Bases: Module

Heterogeneous Graph Transformer (HGT) model for learning on heterogeneous graphs.

This implementation builds on PyTorch Geometric’s HGTConv and supports node-type specific input projection, multi-layer HGT attention, and per-type decoders.

Parameters:

  • in_channels (int) – Input feature dimension for each node.

  • hidden_channels (int) – Hidden feature dimension used in HGT layers.

  • num_heads (int) – Number of attention heads in each HGTConv layer.

  • num_layers (int) – Number of stacked HGTConv layers.

  • n_dec_l (int) – Number of layers in decoder. If 1, uses a single linear layer.

  • data_obj (HeteroData) – A PyG HeteroData object describing the node and edge types in the graph.

  • out_channels (int, optional) – If specified, adds an intermediate projection layer to this dimension before decoding.

lin_dict

Node-type specific input projections to hidden space.

Type:

nn.ModuleDict

convs

List of stacked HGTConv layers.

Type:

nn.ModuleList

decoder

Node-type specific decoders for reconstruction.

Type:

nn.ModuleDict

lin

Optional projection layer if out_channels is specified.

Type:

nn.Linear, optional

node_type

The target node type for downstream prediction.

Type:

str

forward(x_dict, edge_index_dict)[source]

Follow https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/hgt_dblp.py

Forward pass of the HGT model.

Parameters:

  • x_dict (Dict[str, Tensor]) – Dictionary mapping node types to feature matrices.

  • edge_index_dict (Dict[Tuple[str, str, str], Tensor]) – Dictionary mapping edge types to edge index tensors.

Returns:

  • z (Tensor) – The final node embeddings for the primary node type, L2 normalized.

  • recon (Tensor) – Reconstructed input for the primary node type (for autoencoder training).

Classes

spamosaic.architectures.hgt.HGT

Heterogeneous Graph Transformer (HGT) model for learning on heterogeneous graphs.