GCN does not learn

[RituPande]

I attempted to create a GCN model that does not seem to learn. Can someone help me figure out what the issue could be. Maybe there is something basic that is missing my eye.

!pip install numpy
!pip install tensorflow
!pip install spektral
#!pip install tqdm

import numpy as np
import tensorflow as tf
import spektral

cora_dataset = spektral.datasets.citation.Citation(name='cora')
train_mask = cora_dataset.mask_tr
val_mask = cora_dataset.mask_va
test_mask = cora_dataset.mask_te 

cora_graph = cora_dataset.read()[0]

node_features  = cora_graph.x.astype('float32')
node_labels = cora_graph.y.astype('float32')
A = cora_graph.a.astype('float32').toarray()
A = A + np.eye(A.shape[0])
def masked_cross_entropy_loss(  labels,logits, mask ):

  loss = tf.nn.softmax_cross_entropy_with_logits(labels,logits )
  mask = tf.cast(mask, dtype=tf.float32)
  # the step below is important, because we need to find mean of only masked nodes
  # dividing the mask by its mean = mask * total_nodes/total_masked nodes, comes in handy when we try to take the mean of the loss in the final step
  # the total number of nodes that are input to the function are cancelled out between two means of mask and loss
  # What remains is only the total number of masked nodes in denominator.
  mask /= tf.reduce_mean(mask)
  loss *= mask
  return tf.reduce_mean(loss)

def masked_accuracy( labels, logits, mask ):
  accuracy_array = tf.equal(tf.argmax(logits, axis=1), tf.argmax(labels, axis=1))
  accuracy_array = tf.cast(accuracy_array, dtype =tf.float32)
  mask = tf.cast(mask, dtype = tf.float32)
  mask/= tf.reduce_mean(mask)
  accuracy_array *= mask
  return tf.reduce_mean(accuracy_array) 

class GCNLayer:
  def __init__( self, A ):
    self.A = A
       
  
  def _transform( self, units, features, trans_func ):
    if trans_func == 'dense':
        features =  tf.keras.layers.Dense(units)( features )
        features = tf.cast(features, dtype=tf.float32)
        return features 
    else:
        raise Exception('Transformation function not implemented')

  def _aggregate( self, features, agg_func ):
    if agg_func == 'adj_matmul':
      return   self.A @ features 
    else:
      raise Exception('Aggregation function not implemented')

  def _activate( self, features, activation ):
      features =  tf.keras.layers.Activation(activation)( features)
      return features

  def __call__( self, units, features, trans_func='dense', agg_func = 'adj_matmul', activation='relu' ):
    features = self._transform(units, features, trans_func )
    features = self._aggregate(features, agg_func)
    if  activation is not None: : 
      features =  self._activate(features, activation)
    return features

class MyModel:
  def __init__( self, A, node_features, node_labels, train_mask, val_mask, test_mask ):
    self.A = A
    self.node_features = node_features
    self.node_labels = node_labels
    self.train_mask = train_mask
    self.val_mask = val_mask
    self.test_mask = test_mask
    
    self.gcn_layer1 =  GCNLayer(self.A)
    self.gcn_layer2 =  GCNLayer(self.A)

  def __call__( self ):
    hidden_out = self.gcn_layer1(32, self.node_features, activation='relu'  )
    output  = self.gcn_layer2(7, hidden_out, activation=None) 
    return output

  def train( self, num_epochs=1, lr =0.01 ):
    optimizer = tf.keras.optimizers.Adam(lr)
    best_val_acc = 0.0
    for e in range(num_epochs):
      with tf.GradientTape() as t:
        logits = self()
        train_loss = masked_cross_entropy_loss( self.node_labels, logits, self.train_mask )
      
      variables = t.watched_variables()
      grads = t.gradient(train_loss, variables)
      optimizer.apply_gradients(zip(grads, variables))
      logits = self()
      train_acc = masked_accuracy(  self.node_labels, logits, self.train_mask )
      val_acc = masked_accuracy(  self.node_labels, logits, self.val_mask )
      if val_acc > best_val_acc:
        best_val_acc = val_acc
        print(f'epoch={e},Training Loss:{train_loss.numpy()},Training Accuracy:{train_acc.numpy()}, Validation Accuracy:{val_acc.numpy()}')


model = MyModel(A, node_features, node_labels, train_mask, val_mask, test_mask)
model.train(num_epochs=200, lr=0.01)

Output:

epoch=0,Training Loss:7.667728424072266,Training Accuracy:0.12857143580913544, Validation Accuracy:0.1720000058412552
epoch=10,Training Loss:4.182003974914551,Training Accuracy:0.19285714626312256, Validation Accuracy:0.20999999344348907
epoch=17,Training Loss:4.869807720184326,Training Accuracy:0.18571428954601288, Validation Accuracy:0.2279999852180481
epoch=18,Training Loss:7.179490089416504,Training Accuracy:0.15714284777641296, Validation Accuracy:0.25999999046325684
epoch=37,Training Loss:5.826307773590088,Training Accuracy:0.3642857074737549, Validation Accuracy:0.3359999656677246