GSgnnEdgeModelBase

class graphstorm.model.GSgnnEdgeModelBase(*args, **kwargs)

Bases: GSgnnModelBase, GSgnnEdgeModelInterface

GraphStorm GNN model base class for edge prediction tasks.

This base class extends GraphStorm GSgnnModelBase and GSgnnEdgeModelInterface. When users want to define a customized edge prediction GNN model and train the model in GraphStorm, the model class needs to inherit from this base class, and implement the required methods including forward(), predict(), save_model(), restore_model() and create_optimizer().