Building our GCN model with subclassing using Spektral/Keras

PROTEINS_embedding.py

# Now, we can use model subclassing to define our model:

class ProteinsGNN(Model):
  
  def __init__(self, n_hidden, n_labels):
    super().__init__()
    # Define our GCN layer with our n_hidden layers
    self.graph_conv = GCNConv(n_hidden)
    # Define our global pooling layer
    self.pool = GlobalSumPool()
    # Define our dropout layer, initialize dropout freq. to .5 (50%)
    self.dropout = Dropout(0.5)
    # Define our Dense layer, with softmax activation function
    self.dense = Dense(n_labels, 'softmax')

# Define class method to call model on input
  def call(self, inputs):
    out = self.graph_conv(inputs)
    out = self.dropout(out)
    out = self.pool(out)
    out = self.dense(out)

    return out