GraphSAGE Definition

sage_def.py

from dgl.nn.tensorflow import SAGEConv

class SAGE(tf.keras.Model):
    def __init__(self, feat_dim, hidden_dim, class_num):
        super(SAGE, self).__init__()
        self.h1 = SAGEConv(in_feats=feat_dim, out_feats=hidden_dim, aggregator_type='pool', feat_drop=0.5, activation=tf.nn.relu)
        self.h2 = SAGEConv(in_feats=feat_dim, out_feats=class_num, aggregator_type='pool')
    
    def call(self, g, features):
        h = features
        h = self.h1(g, h)
        h = self.h2(g, h)
        return h