GCN using adjacency matrix in sparse COOrdinate format

[RituPande]

I attempted to modify the GCN layer that I created to work with adjacency sparse COOrdinate format. My understanding is that there should be no difference in the output of the model with adjacency matrix in dense or parse COOrdinate format However, I see different outputs. I am pasting my code for the GCN layer that stores the adjacency matrix based on "sparse" parameter to the constructor.

Can you please check if I have correctly implemented the aggregation of nodes using adjacency matrix is sparse COOrdinate format correctly? I have printed the inputs and outputs of all the functions used in the _aggregate function and they seem to be fine to me, unless I have misunderstood the concept somehow.

Then if I print the output of the _aggregate function, I get different values of the aggregated features between sparse and dense formats, though I am not able to find why. I even tried to sort the receiver_nodes and used segment_sum, yet the value of the aggregated features are different from the values of their values when the layer uses dense adjacency matrix

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

import numpy as np
import tensorflow as tf
import spektral
import scipy
from spektral.datasets.citation import Cora
tf.random.set_seed(seed=0)

dataset = Cora(normalize_x=True)
graph = dataset[0]
node_features, A, node_labels = graph.x, graph.a, graph.y
print (type(A))
A = A.astype('float32').toarray()
A = A + np.eye(A.shape[0])
train_mask, val_mask, test_mask = dataset.mask_tr, dataset.mask_va, dataset.mask_te

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, units, activation, sparse = False ):
      
        self.is_sparse = sparse
        if sparse:
          self.a_sender_nodes, self.a_receiver_nodes, _ = scipy.sparse.find(A)
        else:
          self.A = A
        self.kernel = tf.keras.layers.Dense(units)
        self.activation = tf.keras.layers.Activation(activation)
        
    
    def _transform( self, features ):
        features =  self.kernel( features )
        features = tf.cast(features, dtype=tf.float32)
        return features 
        
    def _aggregate( self, features):
        if self.is_sparse:
            X = tf.gather(features, self.a_receiver_nodes )
            unique_segments, _ = tf.unique(self.a_receiver_nodes)
            n_segments =   tf.size(unique_segments)
            features = tf.math.unsorted_segment_sum(data = X, segment_ids = self.a_receiver_nodes, num_segments= n_segments )
            return features
        else:       
            return  self.A @ features 
      
    def _activate( self, features ):
        features =  self.activation( features)
        return features

    def __call__( self, features ):
      features = self._transform( features )
      features = self._aggregate(features)
      features =  self._activate(features)
      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.d0 = tf.keras.layers.Dropout(0.2 )
        self.gcn_layer1 =  GCNLayer(self.A, node_features.shape[1], activation='relu', sparse=True)
        self.gcn_layer2 =  GCNLayer(self.A, node_labels.shape[1], activation = None, sparse=True)
        self.d1 = tf.keras.layers.Dropout(0.2)

    def __call__( self ):
        features = self.d0(self.node_features)
        features =  self.gcn_layer1( features  )
        features = self.d1(features)
        output  = self.gcn_layer2(features) 
        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)

When sparse parameter to GCNLayer is "False", the output is as expected:

epoch=0,Training Loss:1.89188551902771,Training Accuracy:0.3214285671710968, Validation Accuracy:0.24400000274181366
epoch=1,Training Loss:3.431786060333252,Training Accuracy:0.4214285612106323, Validation Accuracy:0.3779999911785126
epoch=2,Training Loss:6.477021217346191,Training Accuracy:0.3642857074737549, Validation Accuracy:0.37199997901916504
epoch=3,Training Loss:7.120337963104248,Training Accuracy:0.5, Validation Accuracy:0.2879999876022339
epoch=4,Training Loss:2.3660573959350586,Training Accuracy:0.7285714149475098, Validation Accuracy:0.44599997997283936
epoch=5,Training Loss:1.2946234941482544,Training Accuracy:0.7214285731315613, Validation Accuracy:0.4819999933242798
epoch=6,Training Loss:1.5609763860702515,Training Accuracy:0.721428632736206, Validation Accuracy:0.47199997305870056
epoch=7,Training Loss:1.660966157913208,Training Accuracy:0.8285714387893677, Validation Accuracy:0.6000000238418579
epoch=8,Training Loss:1.218860149383545,Training Accuracy:0.8428571224212646, Validation Accuracy:0.6739999651908875
epoch=9,Training Loss:0.9399302005767822,Training Accuracy:0.8928571343421936, Validation Accuracy:0.7299999594688416
epoch=10,Training Loss:0.6795181632041931,Training Accuracy:0.8928571343421936, Validation Accuracy:0.7339999675750732
epoch=11,Training Loss:0.5675678253173828,Training Accuracy:0.8785713911056519, Validation Accuracy:0.6959999799728394
epoch=12,Training Loss:0.6819177865982056,Training Accuracy:0.9142857193946838, Validation Accuracy:0.722000002861023
epoch=13,Training Loss:0.5002830624580383,Training Accuracy:0.9571428894996643, Validation Accuracy:0.7219999432563782
epoch=14,Training Loss:0.3572576940059662,Training Accuracy:0.9571428894996643, Validation Accuracy:0.6940000057220459
epoch=15,Training Loss:0.38455110788345337,Training Accuracy:0.9500000476837158, Validation Accuracy:0.6660000085830688
epoch=16,Training Loss:0.40155625343322754,Training Accuracy:0.9642857909202576, Validation Accuracy:0.6859999895095825
epoch=17,Training Loss:0.3217295706272125,Training Accuracy:0.971428632736206, Validation Accuracy:0.699999988079071
epoch=18,Training Loss:0.22483675181865692,Training Accuracy:0.9785715341567993, Validation Accuracy:0.722000002861023
epoch=19,Training Loss:0.20372000336647034,Training Accuracy:0.9785715341567993, Validation Accuracy:0.7280000448226929
epoch=20,Training Loss:0.17791800200939178,Training Accuracy:0.9857143759727478, Validation Accuracy:0.75
epoch=21,Training Loss:0.12488004565238953,Training Accuracy:0.9785715341567993, Validation Accuracy:0.7620000243186951
epoch=22,Training Loss:0.10189366340637207,Training Accuracy:0.9857143759727478, Validation Accuracy:0.7700000405311584
epoch=23,Training Loss:0.08439365774393082,Training Accuracy:1.0, Validation Accuracy:0.768000066280365
epoch=24,Training Loss:0.056988950818777084,Training Accuracy:1.0, Validation Accuracy:0.7619999647140503
epoch=25,Training Loss:0.049707140773534775,Training Accuracy:1.0, Validation Accuracy:0.7639999389648438
epoch=26,Training Loss:0.04378794506192207,Training Accuracy:1.0, Validation Accuracy:0.765999972820282
epoch=27,Training Loss:0.03783036768436432,Training Accuracy:1.0, Validation Accuracy:0.7679999470710754
epoch=28,Training Loss:0.032043490558862686,Training Accuracy:1.0, Validation Accuracy:0.7600000500679016
epoch=29,Training Loss:0.02692755125463009,Training Accuracy:1.0, Validation Accuracy:0.7580000162124634
epoch=30,Training Loss:0.022733008489012718,Training Accuracy:1.0, Validation Accuracy:0.7660000324249268
epoch=31,Training Loss:0.019153233617544174,Training Accuracy:1.0, Validation Accuracy:0.7699999809265137
epoch=32,Training Loss:0.015598902478814125,Training Accuracy:1.0, Validation Accuracy:0.7640000581741333
epoch=33,Training Loss:0.012799952179193497,Training Accuracy:1.0, Validation Accuracy:0.7660000324249268
epoch=34,Training Loss:0.010552829131484032,Training Accuracy:1.0, Validation Accuracy:0.768000066280365
epoch=35,Training Loss:0.008755798451602459,Training Accuracy:1.0, Validation Accuracy:0.7660000324249268

When the sparse is "True", the validation accuracy seems to worsen with time, and training accuracy increases too rapidly ( when I expected it to be same as above )

epoch=0,Training Loss:2.1233248710632324,Training Accuracy:0.15714284777641296, Validation Accuracy:0.12199999392032623
epoch=1,Training Loss:3.4777023792266846,Training Accuracy:0.5714285969734192, Validation Accuracy:0.11999998241662979
epoch=2,Training Loss:1.3093963861465454,Training Accuracy:0.8142856955528259, Validation Accuracy:0.24400000274181366
epoch=3,Training Loss:1.1654517650604248,Training Accuracy:0.8785713911056519, Validation Accuracy:0.4259999692440033
epoch=4,Training Loss:0.5752173066139221,Training Accuracy:0.9000000953674316, Validation Accuracy:0.5020000338554382
epoch=5,Training Loss:0.4111606776714325,Training Accuracy:0.9428572058677673, Validation Accuracy:0.4819999933242798
epoch=6,Training Loss:0.3007441759109497,Training Accuracy:0.9857143759727478, Validation Accuracy:0.4560000002384186
epoch=7,Training Loss:0.21603146195411682,Training Accuracy:1.0, Validation Accuracy:0.4519999921321869
epoch=8,Training Loss:0.1521170288324356,Training Accuracy:1.0, Validation Accuracy:0.47200000286102295
epoch=9,Training Loss:0.10611247271299362,Training Accuracy:1.0, Validation Accuracy:0.48799997568130493
epoch=10,Training Loss:0.07399296760559082,Training Accuracy:1.0, Validation Accuracy:0.4819999933242798
epoch=11,Training Loss:0.0518457256257534,Training Accuracy:1.0, Validation Accuracy:0.4439999759197235
epoch=12,Training Loss:0.03679153323173523,Training Accuracy:1.0, Validation Accuracy:0.4059999883174896
epoch=13,Training Loss:0.02675839327275753,Training Accuracy:1.0, Validation Accuracy:0.3580000102519989
epoch=14,Training Loss:0.019282184541225433,Training Accuracy:1.0, Validation Accuracy:0.32199999690055847
epoch=15,Training Loss:0.013474024832248688,Training Accuracy:1.0, Validation Accuracy:0.3140000104904175
epoch=16,Training Loss:0.009401041083037853,Training Accuracy:1.0, Validation Accuracy:0.3059999942779541
epoch=17,Training Loss:0.00659525953233242,Training Accuracy:1.0, Validation Accuracy:0.29600000381469727
epoch=18,Training Loss:0.0046736085787415504,Training Accuracy:1.0, Validation Accuracy:0.2919999957084656
epoch=19,Training Loss:0.00335993361659348,Training Accuracy:1.0, Validation Accuracy:0.2919999957084656
epoch=20,Training Loss:0.0024560813326388597,Training Accuracy:1.0, Validation Accuracy:0.28599998354911804
epoch=21,Training Loss:0.0018257534829899669,Training Accuracy:1.0, Validation Accuracy:0.2840000092983246
epoch=22,Training Loss:0.0013794495025649667,Training Accuracy:1.0, Validation Accuracy:0.28200000524520874
epoch=23,Training Loss:0.0010587455471977592,Training Accuracy:1.0, Validation Accuracy:0.2839999794960022
epoch=24,Training Loss:0.000825054245069623,Training Accuracy:1.0, Validation Accuracy:0.28200000524520874
epoch=25,Training Loss:0.0006526490324176848,Training Accuracy:1.0, Validation Accuracy:0.2800000011920929
epoch=26,Training Loss:0.0005238134181126952,Training Accuracy:1.0, Validation Accuracy:0.28200000524520874
epoch=27,Training Loss:0.0004265294410288334,Training Accuracy:1.0, Validation Accuracy:0.2840000092983246
epoch=28,Training Loss:0.00035230693174526095,Training Accuracy:1.0, Validation Accuracy:0.2800000011920929
epoch=29,Training Loss:0.0002949906629510224,Training Accuracy:1.0, Validation Accuracy:0.28200000524520874
epoch=30,Training Loss:0.0002502320858184248,Training Accuracy:1.0, Validation Accuracy:0.2800000011920929
epoch=31,Training Loss:0.00021486477635335177,Training Accuracy:1.0, Validation Accuracy:0.27799999713897705
epoch=32,Training Loss:0.00018666293181013316,Training Accuracy:1.0, Validation Accuracy:0.2759999930858612
epoch=33,Training Loss:0.0001639004040043801,Training Accuracy:1.0, Validation Accuracy:0.27799999713897705
epoch=34,Training Loss:0.00014533725334331393,Training Accuracy:1.0, Validation Accuracy:0.2759999930858612
epoch=35,Training Loss:0.00013006437802687287,Training Accuracy:1.0, Validation Accuracy:0.2759999930858612

[RituPande]

I found the the issue. The segment_sum had to be done on receiver nodes while the tf_gather had to be done on sendernodes. I was doing both on the receiver nodes which was causing the issue. The issue is now resolved.