graphstorm.model

A GraphStorm model may contain three components:

  • Input layer: a set of modules to convert input data for different use cases, e.g., embedding texture features.

  • Encoder: a set of Graph Neural Network modules

  • Decoder: a set of modules to convert results from encoders for different tasks, e.g., classification, regression, or link prediction.

Currently GraphStorm releases the first two set of components.

If users would like to implement their own model, the best practice is to extend the corresponding ***ModelBase, and implement the abstract methods.

Base models

GSgnnNodeModelBase

The base class for node-prediction GNN

GSgnnEdgeModelBase

The base class for edge-prediction GNN

GSgnnLinkPredictionModelBase

The base class for link-prediction GNN

Input Layers

GSNodeEncoderInputLayer

The input encoder layer for all nodes in a heterogeneous graph.

GSLMNodeEncoderInputLayer

The input encoder layer with language model for all nodes in a heterogeneous graph.

GSPureLMNodeInputLayer

The input embedding layer with language model only for all nodes in a heterogeneous graph.

Encoders and GNN Layers

RelationalGCNEncoder

Relational graph conv encoder.

RelGraphConvLayer

Relational graph convolution layer from Modeling Relational Data with Graph Convolutional Networks.

RelationalGATEncoder

Relational graph attention encoder

RelationalAttLayer

Relational graph attention layer from ` Relational Graph Attention Networks <https://arxiv.org/abs/1904.05811>`__.

SAGEEncoder

GraphSage Conv Encoder

SAGEConv

GraphSage Convolutional layerfrom Inductive Representation Learning on Large Graphs

HGTEncoder

Heterogenous graph transformer (HGT) encoder

HGTLayer

Heterogenous graph transformer (HGT) layer from Heterogeneous Graph Transformer.