Pass keyword arguments to class init, Modify DEFINE_flags

Author: joshchang1112

research/gnn-survey-paper/layers.py

@@ -1,29 +1,34 @@
 import tensorflow as tf
 
 class GraphConvLayer(tf.keras.layers.Layer):
+    """ Single graph convolution layer
 
-    def __init__(self, features):
+    Args:
+        input_dim (int): Input dimension of gcn layer
+        output_dim (int): Output dimension of gcn layer
+        bias (bool): Whether bias needs to be added to the layer
+    """
+    def __init__(self, **kwargs):
         super(GraphConvLayer, self).__init__()
-        self.in_feat = features['input_dim']
-        self.out_feat = features['output_dim']
-        self.b = features['bias']
-
+        for key, item in kwargs.items():
+            setattr(self, key, item)
+        
     def build(self, input_shape):
         self.weight = self.add_weight(name="weight",
-                                      shape=(self.in_feat, self.out_feat),
+                                      shape=(self.input_dim, self.output_dim),
                                       initializer='random_normal',
                                       trainable=True)
-        if self.b:
-            self.bias = self.add_weight(name="bias",
-                                        shape=(self.out_feat,),
+        if self.bias:
+            self.b = self.add_weight(name="bias",
+                                        shape=(self.output_dim,),
                                         initializer='random_normal',
                                         trainable=True)
     def call(self, inputs):
         x, adj = inputs[0], inputs[1]
-       	x = tf.matmul(adj, x)
+        x = tf.matmul(adj, x)
         outputs = tf.matmul(x, self.weight)
-        if self.b:
-            return self.bias + outputs
+        if self.bias:
+            return self.b + outputs
         else:
             return outputs
 

research/gnn-survey-paper/models.py

@@ -2,37 +2,47 @@
 from layers import GraphConvLayer
 
 class GCN(tf.keras.Model):
-
-    def __init__(self, features_dim, num_layers, hidden_dim, num_classes, dropout_rate, bias=True):
+    """Graph convolution network for semi-supevised node classification.
+
+    Args:
+        features_dim (int): Dimension of input features 
+        num_layers (int): Number of gnn layers 
+        hidden_dim (list): List of hidden layers dimension
+        num_classes (int): Total number of classes
+        dropout_prob (float): Dropout probability
+        bias (bool): Whether bias needs to be added to gcn layers
+    """
+    
+    def __init__(self, **kwargs):
         super(GCN, self).__init__()
-
-        self.num_layers = num_layers
-        self.bias = bias
-
+        
+        for key, item in kwargs.items():
+            setattr(self, key, item)
+        
         self.gc = []
         # input layer
         single_gc = tf.keras.Sequential()
-        single_gc.add(GraphConvLayer({"input_dim": features_dim,
-                                      "output_dim": hidden_dim[0],
-                                      "bias": bias}))
+        single_gc.add(GraphConvLayer(input_dim=self.features_dim,
+                                     output_dim=self.hidden_dim[0],
+                                     bias=self.bias))
         single_gc.add(tf.keras.layers.ReLU())
-        single_gc.add(tf.keras.layers.Dropout(dropout_rate))
+        single_gc.add(tf.keras.layers.Dropout(self.dropout_prob))
         self.gc.append(single_gc)
 
         # hidden layers
-        for i in range(0, num_layers-2):
+        for i in range(0, self.num_layers-2):
             single_gc = tf.keras.Sequential()
-            single_gc.add(GraphConvLayer({"input_dim": hidden_dim[i],
-                                          "output_dim": hidden_dim[i+1],
-                                          "bias": bias}))
+            single_gc.add(GraphConvLayer(input_dim=self.hidden_dim[i],
+                                         output_dim=self.hidden_dim[i+1],
+                                         bias=self.bias))
             single_gc.add(tf.keras.layers.ReLU())
-            single_gc.add(tf.keras.layers.Dropout(dropout_rate))
+            single_gc.add(tf.keras.layers.Dropout(self.dropout_prob))
             self.gc.append(single_gc)
 
         # output layer
-        self.classifier = GraphConvLayer({"input_dim": hidden_dim[-1],
-                                          "output_dim": num_classes,
-                                          "bias": bias})
+        self.classifier = GraphConvLayer(input_dim=self.hidden_dim[-1],
+                                         output_dim=self.num_classes,
+                                         bias=self.bias)
 
     def call(self, inputs):
         features, adj = inputs[0], inputs[1]

research/gnn-survey-paper/train.py

@@ -8,17 +8,17 @@
 
 from utils import load_dataset, build_model, cal_acc
 
-flags.DEFINE_string('dataset', 'cora',
+flags.DEFINE_enum('dataset', 'cora', ['cora'],
                     'The input dataset. Avaliable dataset now: cora')
-flags.DEFINE_string('model', 'gcn',
-                    'GNN model. Available model now: gcn')
+flags.DEFINE_enum('model', 'gcn', ['gcn'],
+                  'GNN model. Available model now: gcn')
 flags.DEFINE_float('dropout', 0.5, 'Dropout probability')
 flags.DEFINE_integer('gpu', '-1', 'Gpu id, -1 means cpu only')
 flags.DEFINE_float('lr', 1e-2, 'Initial learning rate')
 flags.DEFINE_integer('epochs', 200, 'Number of training epochs')
 flags.DEFINE_integer('num_layers', 2, 'Number of gnn layers')
 flags.DEFINE_list('hidden_dim', [32], 'Dimension of gnn hidden layers')
-flags.DEFINE_string('optimizer', 'adam', 'Optimizer for training')
+flags.DEFINE_enum('optimizer', 'adam', ['adam', 'sgd'], 'Optimizer for training')
 flags.DEFINE_float('weight_decay', 5e-4, 'Weight for L2 regularization')
 flags.DEFINE_string('save_dir', 'models/cora/gcn', 'Directory stores trained model')
 
@@ -33,6 +33,7 @@ def train(model, adj, features, labels, idx_train, idx_val, idx_test):
         optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.lr)
     elif FLAGS.optimizer == 'sgd':
         optimizer = tf.keras.optimizers.SGD(learning_rate=FLAGS.lr)
+    
 
     inputs = (features, adj)
     for epoch in range(FLAGS.epochs):
@@ -70,8 +71,6 @@ def train(model, adj, features, labels, idx_train, idx_val, idx_test):
     print("***Test Accuracy: %.3f***"% (test_acc))
 
 
-
-
 def main(_):
 
     if FLAGS.gpu == -1:

research/gnn-survey-paper/utils.py

@@ -13,7 +13,12 @@ def build_model(model_name, features_dim, num_layers, hidden_dim, num_classes, d
 
     # Only gcn available now
     if model_name == 'gcn':
-        model = GCN(features_dim, num_layers, hidden_dim, num_classes, dropout)
+        model = GCN(features_dim=features_dim, 
+                    num_layers=num_layers, 
+                    hidden_dim=hidden_dim, 
+                    num_classes=num_classes, 
+                    dropout_prob=dropout,
+                    bias=True)
 
     elif model_name == 'gat':
         raise NotImplementedError