diff --git a/src/openne/__main__.py b/src/openne/__main__.py index e6907492..aaa47b17 100644 --- a/src/openne/__main__.py +++ b/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) diff --git a/src/openne/gcn/gcnAPI.py b/src/openne/gcn/gcnAPI.py index fa58813c..79d2b6f9 100644 --- a/src/openne/gcn/gcnAPI.py +++ b/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): diff --git a/src/openne/gcn/layers.py b/src/openne/gcn/layers.py index a923fb09..c0003dbd 100644 --- a/src/openne/gcn/layers.py +++ b/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)