save embedding vectors of GCN

src/openne/__main__.py

@@ -163,9 +163,8 @@ def main(args):
                                       epoch=args.epochs, learning_rate=args.lr, weight_decay=args.weight_decay)
     t2 = time.time()
     print(t2-t1)
-    if args.method != 'gcn':
-        print("Saving embeddings...")
-        model.save_embeddings(args.output)
+    print("Saving embeddings...")
+    model.save_embeddings(args.output)
     if args.label_file and args.method != 'gcn':
         vectors = model.vectors
         X, Y = read_node_label(args.label_file)

src/openne/gcn/gcnAPI.py

@@ -40,6 +40,8 @@ def __init__(self, graph, learning_rate=0.01, epochs=200,
         # Init variables
         self.sess.run(tf.global_variables_initializer())
 
+        self.vectors = {}
+
         cost_val = []
 
         # Train model
@@ -51,17 +53,17 @@ def __init__(self, graph, learning_rate=0.01, epochs=200,
             feed_dict.update({self.placeholders['dropout']: self.dropout})
 
             # Training step
-            outs = self.sess.run(
-                [self.model.opt_op, self.model.loss, self.model.accuracy], feed_dict=feed_dict)
+            self.outs = self.sess.run(
+                [self.model.opt_op, self.model.loss, self.model.accuracy, self.model.layers[0].embedding], feed_dict=feed_dict)
 
             # Validation
             cost, acc, duration = self.evaluate(self.val_mask)
             cost_val.append(cost)
 
             # Print results
-            print("Epoch:", '%04d' % (epoch + 1), "train_loss=", "{:.5f}".format(outs[1]),
+            print("Epoch:", '%04d' % (epoch + 1), "train_loss=", "{:.5f}".format(self.outs[1]),
                   "train_acc=", "{:.5f}".format(
-                outs[2]), "val_loss=", "{:.5f}".format(cost),
+                self.outs[2]), "val_loss=", "{:.5f}".format(cost),
                 "val_acc=", "{:.5f}".format(acc), "time=", "{:.5f}".format(time.time() - t))
 
             if epoch > self.early_stopping and cost_val[-1] > np.mean(cost_val[-(self.early_stopping+1):-1]):
@@ -74,6 +76,19 @@ def __init__(self, graph, learning_rate=0.01, epochs=200,
         print("Test set results:", "cost=", "{:.5f}".format(test_cost),
               "accuracy=", "{:.5f}".format(test_acc), "time=", "{:.5f}".format(test_duration))
 
+        self.embeddings = self.outs[3]
+        for i, embedding in enumerate(self.embeddings):
+            self.vectors[i] = embedding
+
+    # save embedding function
+    def save_embeddings(self, filename):
+        fout = open(filename, 'w')
+        node_num = len(self.vectors)
+        fout.write("{} {}\n".format(node_num, self.hidden1))
+        for node, vec in self.vectors.items():
+            fout.write("{} {}\n".format(node, ' '.join([str(x) for x in vec])))
+        fout.close()
+
     # Define model evaluation function
 
     def evaluate(self, mask):

src/openne/gcn/layers.py

@@ -186,5 +186,5 @@ def _call(self, inputs):
         # bias
         if self.bias:
             output += self.vars['bias']
-
+        self.embedding = output
         return self.act(output)