Skip to content

GitLab

Unverified Commit be8763fa authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Black auto fix. (#4679) 


Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent eae6ce2a
Hide whitespace changes
Inline Side-by-side
Showing with 413 additions and 277 deletions
examples/tensorflow/gcn/gcn_mp.py
View file @ be8763fa
import argparse import argparse
import time
import math import math
import numpy as np import time
import networkx as nx import networkx as nx
import numpy as np
import tensorflow as tf import tensorflow as tf
import dgl
from dgl.data import register_data_args
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
from tensorflow.keras import layers from tensorflow.keras import layers
import dgl
from dgl.data import (
CiteseerGraphDataset,
CoraGraphDataset,
PubmedGraphDataset,
register_data_args,
)
def gcn_msg(edge): def gcn_msg(edge):
msg = edge.src['h'] * edge.src['norm'] msg = edge.src["h"] * edge.src["norm"]
return {'m': msg} return {"m": msg}
def gcn_reduce(node): def gcn_reduce(node):
accum = tf.reduce_sum(node.mailbox['m'], 1) * node.data['norm'] accum = tf.reduce_sum(node.mailbox["m"], 1) * node.data["norm"]
return {'h': accum} return {"h": accum}
class GCNLayer(layers.Layer): class GCNLayer(layers.Layer):
def __init__(self, def __init__(self, g, in_feats, out_feats, activation, dropout, bias=True):
g,
in_feats,
out_feats,
activation,
dropout,
bias=True):
super(GCNLayer, self).__init__() super(GCNLayer, self).__init__()
self.g = g self.g = g
w_init = tf.random_normal_initializer() w_init = tf.random_normal_initializer()
self.weight = tf.Variable(initial_value=w_init(shape=(in_feats, out_feats), self.weight = tf.Variable(
dtype='float32'), initial_value=w_init(shape=(in_feats, out_feats), dtype="float32"),
trainable=True) trainable=True,
)
if dropout: if dropout:
self.dropout = layers.Dropout(rate=dropout) self.dropout = layers.Dropout(rate=dropout)
else: else:
self.dropout = 0. self.dropout = 0.0
if bias: if bias:
b_init = tf.zeros_initializer() b_init = tf.zeros_initializer()
self.bias = tf.Variable(initial_value=b_init(shape=(out_feats,), self.bias = tf.Variable(
dtype='float32'), initial_value=b_init(shape=(out_feats,), dtype="float32"),
trainable=True) trainable=True,
)
else: else:
self.bias = None self.bias = None
self.activation = activation self.activation = activation
...@@ -50,9 +53,9 @@ class GCNLayer(layers.Layer): ...@@ -50,9 +53,9 @@ class GCNLayer(layers.Layer):
def call(self, h): def call(self, h):
if self.dropout: if self.dropout:
h = self.dropout(h) h = self.dropout(h)
self.g.ndata['h'] = tf.matmul(h, self.weight) self.g.ndata["h"] = tf.matmul(h, self.weight)
self.g.update_all(gcn_msg, gcn_reduce) self.g.update_all(gcn_msg, gcn_reduce)
h = self.g.ndata['h'] h = self.g.ndata["h"]
if self.bias is not None: if self.bias is not None:
h = h + self.bias h = h + self.bias
if self.activation: if self.activation:
...@@ -61,24 +64,19 @@ class GCNLayer(layers.Layer): ...@@ -61,24 +64,19 @@ class GCNLayer(layers.Layer):
class GCN(layers.Layer): class GCN(layers.Layer):
def __init__(self, def __init__(
g, self, g, in_feats, n_hidden, n_classes, n_layers, activation, dropout
in_feats, ):
n_hidden,
n_classes,
n_layers,
activation,
dropout):
super(GCN, self).__init__() super(GCN, self).__init__()
self.layers = [] self.layers = []
# input layer # input layer
self.layers.append( self.layers.append(GCNLayer(g, in_feats, n_hidden, activation, dropout))
GCNLayer(g, in_feats, n_hidden, activation, dropout))
# hidden layers # hidden layers
for i in range(n_layers - 1): for i in range(n_layers - 1):
self.layers.append( self.layers.append(
GCNLayer(g, n_hidden, n_hidden, activation, dropout)) GCNLayer(g, n_hidden, n_hidden, activation, dropout)
)
# output layer # output layer
self.layers.append(GCNLayer(g, n_hidden, n_classes, None, dropout)) self.layers.append(GCNLayer(g, n_hidden, n_classes, None, dropout))
...@@ -100,14 +98,14 @@ def evaluate(model, features, labels, mask): ...@@ -100,14 +98,14 @@ def evaluate(model, features, labels, mask):
def main(args): def main(args):
# load and preprocess dataset # load and preprocess dataset
if args.dataset == 'cora': if args.dataset == "cora":
data = CoraGraphDataset() data = CoraGraphDataset()
elif args.dataset == 'citeseer': elif args.dataset == "citeseer":
data = CiteseerGraphDataset() data = CiteseerGraphDataset()
elif args.dataset == 'pubmed': elif args.dataset == "pubmed":
data = PubmedGraphDataset() data = PubmedGraphDataset()
else: else:
raise ValueError('Unknown dataset: {}'.format(args.dataset)) raise ValueError("Unknown dataset: {}".format(args.dataset))
g = data[0] g = data[0]
if args.gpu < 0: if args.gpu < 0:
...@@ -117,24 +115,29 @@ def main(args): ...@@ -117,24 +115,29 @@ def main(args):
g = g.to(device) g = g.to(device)
with tf.device(device): with tf.device(device):
features = g.ndata['feat'] features = g.ndata["feat"]
labels = g.ndata['label'] labels = g.ndata["label"]
train_mask = g.ndata['train_mask'] train_mask = g.ndata["train_mask"]
val_mask = g.ndata['val_mask'] val_mask = g.ndata["val_mask"]
test_mask = g.ndata['test_mask'] test_mask = g.ndata["test_mask"]
in_feats = features.shape[1] in_feats = features.shape[1]
n_classes = data.num_labels n_classes = data.num_labels
n_edges = data.graph.number_of_edges() n_edges = data.graph.number_of_edges()
print("""----Data statistics------' print(
"""----Data statistics------'
#Edges %d #Edges %d
#Classes %d #Classes %d
#Train samples %d #Train samples %d
#Val samples %d #Val samples %d
#Test samples %d""" % #Test samples %d"""
(n_edges, n_classes, % (
train_mask.numpy().sum(), n_edges,
val_mask.numpy().sum(), n_classes,
test_mask.numpy().sum())) train_mask.numpy().sum(),
val_mask.numpy().sum(),
test_mask.numpy().sum(),
)
)
# add self loop # add self loop
if args.self_loop: if args.self_loop:
...@@ -147,22 +150,24 @@ def main(args): ...@@ -147,22 +150,24 @@ def main(args):
norm = tf.math.pow(degs, -0.5) norm = tf.math.pow(degs, -0.5)
norm = tf.where(tf.math.is_inf(norm), tf.zeros_like(norm), norm) norm = tf.where(tf.math.is_inf(norm), tf.zeros_like(norm), norm)
g.ndata['norm'] = tf.expand_dims(norm, -1) g.ndata["norm"] = tf.expand_dims(norm, -1)
# create GCN model # create GCN model
model = GCN(g, model = GCN(
in_feats, g,
args.n_hidden, in_feats,
n_classes, args.n_hidden,
args.n_layers, n_classes,
tf.nn.relu, args.n_layers,
args.dropout) tf.nn.relu,
args.dropout,
optimizer = tf.keras.optimizers.Adam( )
learning_rate=args.lr)
optimizer = tf.keras.optimizers.Adam(learning_rate=args.lr)
loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy( loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True) from_logits=True
)
# initialize graph # initialize graph
dur = [] dur = []
for epoch in range(args.n_epochs): for epoch in range(args.n_epochs):
...@@ -177,8 +182,9 @@ def main(args): ...@@ -177,8 +182,9 @@ def main(args):
# of Adam(W) optimizer with PyTorch. And this results in worse results. # of Adam(W) optimizer with PyTorch. And this results in worse results.
# Manually adding weights to the loss to do weight decay solves this problem. # Manually adding weights to the loss to do weight decay solves this problem.
for weight in model.trainable_weights: for weight in model.trainable_weights:
loss_value = loss_value + \ loss_value = loss_value + args.weight_decay * tf.nn.l2_loss(
args.weight_decay*tf.nn.l2_loss(weight) weight
)
grads = tape.gradient(loss_value, model.trainable_weights) grads = tape.gradient(loss_value, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights)) optimizer.apply_gradients(zip(grads, model.trainable_weights))
...@@ -186,33 +192,46 @@ def main(args): ...@@ -186,33 +192,46 @@ def main(args):
dur.append(time.time() - t0) dur.append(time.time() - t0)
acc = evaluate(model, features, labels, val_mask) acc = evaluate(model, features, labels, val_mask)
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | " print(
"ETputs(KTEPS) {:.2f}". format(epoch, np.mean(dur), loss_value.numpy().item(), "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
acc, n_edges / np.mean(dur) / 1000)) "ETputs(KTEPS) {:.2f}".format(
epoch,
np.mean(dur),
loss_value.numpy().item(),
acc,
n_edges / np.mean(dur) / 1000,
)
)
acc = evaluate(model, features, labels, test_mask) acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc)) print("Test Accuracy {:.4f}".format(acc))
if __name__ == '__main__': if __name__ == "__main__":
parser = argparse.ArgumentParser(description='GCN') parser = argparse.ArgumentParser(description="GCN")
register_data_args(parser) register_data_args(parser)
parser.add_argument("--dropout", type=float, default=0.5, parser.add_argument(
help="dropout probability") "--dropout", type=float, default=0.5, help="dropout probability"
parser.add_argument("--gpu", type=int, default=-1, )
help="gpu") parser.add_argument("--gpu", type=int, default=-1, help="gpu")
parser.add_argument("--lr", type=float, default=1e-2, parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
help="learning rate") parser.add_argument(
parser.add_argument("--n-epochs", type=int, default=200, "--n-epochs", type=int, default=200, help="number of training epochs"
help="number of training epochs") )
parser.add_argument("--n-hidden", type=int, default=16, parser.add_argument(
help="number of hidden gcn units") "--n-hidden", type=int, default=16, help="number of hidden gcn units"
parser.add_argument("--n-layers", type=int, default=1, )
help="number of hidden gcn layers") parser.add_argument(
parser.add_argument("--weight-decay", type=float, default=5e-4, "--n-layers", type=int, default=1, help="number of hidden gcn layers"
help="Weight for L2 loss") )
parser.add_argument("--self-loop", action='store_true', parser.add_argument(
help="graph self-loop (default=False)") "--weight-decay", type=float, default=5e-4, help="Weight for L2 loss"
)
parser.add_argument(
"--self-loop",
action="store_true",
help="graph self-loop (default=False)",
)
args = parser.parse_args() args = parser.parse_args()
print(args) print(args)
......
examples/tensorflow/gcn/train.py
View file @ be8763fa
import argparse import argparse
import time import time
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
import dgl
from dgl.data import CoraGraphDataset, CiteseerGraphDataset, PubmedGraphDataset
from gcn import GCN from gcn import GCN
import dgl
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
def evaluate(model, features, labels, mask): def evaluate(model, features, labels, mask):
logits = model(features, training=False) logits = model(features, training=False)
...@@ -18,14 +20,14 @@ def evaluate(model, features, labels, mask): ...@@ -18,14 +20,14 @@ def evaluate(model, features, labels, mask):
def main(args): def main(args):
# load and preprocess dataset # load and preprocess dataset
if args.dataset == 'cora': if args.dataset == "cora":
data = CoraGraphDataset() data = CoraGraphDataset()
elif args.dataset == 'citeseer': elif args.dataset == "citeseer":
data = CiteseerGraphDataset() data = CiteseerGraphDataset()
elif args.dataset == 'pubmed': elif args.dataset == "pubmed":
data = PubmedGraphDataset() data = PubmedGraphDataset()
else: else:
raise ValueError('Unknown dataset: {}'.format(args.dataset)) raise ValueError("Unknown dataset: {}".format(args.dataset))
g = data[0] g = data[0]
if args.gpu < 0: if args.gpu < 0:
...@@ -35,24 +37,29 @@ def main(args): ...@@ -35,24 +37,29 @@ def main(args):
g = g.to(device) g = g.to(device)
with tf.device(device): with tf.device(device):
features = g.ndata['feat'] features = g.ndata["feat"]
labels = g.ndata['label'] labels = g.ndata["label"]
train_mask = g.ndata['train_mask'] train_mask = g.ndata["train_mask"]
val_mask = g.ndata['val_mask'] val_mask = g.ndata["val_mask"]
test_mask = g.ndata['test_mask'] test_mask = g.ndata["test_mask"]
in_feats = features.shape[1] in_feats = features.shape[1]
n_classes = data.num_labels n_classes = data.num_labels
n_edges = data.graph.number_of_edges() n_edges = data.graph.number_of_edges()
print("""----Data statistics------' print(
"""----Data statistics------'
#Edges %d #Edges %d
#Classes %d #Classes %d
#Train samples %d #Train samples %d
#Val samples %d #Val samples %d
#Test samples %d""" % #Test samples %d"""
(n_edges, n_classes, % (
train_mask.numpy().sum(), n_edges,
val_mask.numpy().sum(), n_classes,
test_mask.numpy().sum())) train_mask.numpy().sum(),
val_mask.numpy().sum(),
test_mask.numpy().sum(),
)
)
# add self loop # add self loop
if args.self_loop: if args.self_loop:
...@@ -64,22 +71,26 @@ def main(args): ...@@ -64,22 +71,26 @@ def main(args):
norm = tf.math.pow(degs, -0.5) norm = tf.math.pow(degs, -0.5)
norm = tf.where(tf.math.is_inf(norm), tf.zeros_like(norm), norm) norm = tf.where(tf.math.is_inf(norm), tf.zeros_like(norm), norm)
g.ndata['norm'] = tf.expand_dims(norm, -1) g.ndata["norm"] = tf.expand_dims(norm, -1)
# create GCN model # create GCN model
model = GCN(g, model = GCN(
in_feats, g,
args.n_hidden, in_feats,
n_classes, args.n_hidden,
args.n_layers, n_classes,
tf.nn.relu, args.n_layers,
args.dropout) tf.nn.relu,
args.dropout,
)
loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy( loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True) from_logits=True
)
# use optimizer # use optimizer
optimizer = tf.keras.optimizers.Adam( optimizer = tf.keras.optimizers.Adam(
learning_rate=args.lr, epsilon=1e-8) learning_rate=args.lr, epsilon=1e-8
)
# initialize graph # initialize graph
dur = [] dur = []
...@@ -95,8 +106,9 @@ def main(args): ...@@ -95,8 +106,9 @@ def main(args):
# of Adam(W) optimizer with PyTorch. And this results in worse results. # of Adam(W) optimizer with PyTorch. And this results in worse results.
# Manually adding weights to the loss to do weight decay solves this problem. # Manually adding weights to the loss to do weight decay solves this problem.
for weight in model.trainable_weights: for weight in model.trainable_weights:
loss_value = loss_value + \ loss_value = loss_value + args.weight_decay * tf.nn.l2_loss(
args.weight_decay*tf.nn.l2_loss(weight) weight
)
grads = tape.gradient(loss_value, model.trainable_weights) grads = tape.gradient(loss_value, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights)) optimizer.apply_gradients(zip(grads, model.trainable_weights))
...@@ -104,34 +116,51 @@ def main(args): ...@@ -104,34 +116,51 @@ def main(args):
dur.append(time.time() - t0) dur.append(time.time() - t0)
acc = evaluate(model, features, labels, val_mask) acc = evaluate(model, features, labels, val_mask)
print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | " print(
"ETputs(KTEPS) {:.2f}". format(epoch, np.mean(dur), loss_value.numpy().item(), "Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | Accuracy {:.4f} | "
acc, n_edges / np.mean(dur) / 1000)) "ETputs(KTEPS) {:.2f}".format(
epoch,
np.mean(dur),
loss_value.numpy().item(),
acc,
n_edges / np.mean(dur) / 1000,
)
)
acc = evaluate(model, features, labels, test_mask) acc = evaluate(model, features, labels, test_mask)
print("Test Accuracy {:.4f}".format(acc)) print("Test Accuracy {:.4f}".format(acc))
if __name__ == '__main__': if __name__ == "__main__":
parser = argparse.ArgumentParser(description='GCN') parser = argparse.ArgumentParser(description="GCN")
parser.add_argument("--dataset", type=str, default="cora", parser.add_argument(
help="Dataset name ('cora', 'citeseer', 'pubmed').") "--dataset",
parser.add_argument("--dropout", type=float, default=0.5, type=str,
help="dropout probability") default="cora",
parser.add_argument("--gpu", type=int, default=-1, help="Dataset name ('cora', 'citeseer', 'pubmed').",
help="gpu") )
parser.add_argument("--lr", type=float, default=1e-2, parser.add_argument(
help="learning rate") "--dropout", type=float, default=0.5, help="dropout probability"
parser.add_argument("--n-epochs", type=int, default=200, )
help="number of training epochs") parser.add_argument("--gpu", type=int, default=-1, help="gpu")
parser.add_argument("--n-hidden", type=int, default=16, parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
help="number of hidden gcn units") parser.add_argument(
parser.add_argument("--n-layers", type=int, default=1, "--n-epochs", type=int, default=200, help="number of training epochs"
help="number of hidden gcn layers") )
parser.add_argument("--weight-decay", type=float, default=5e-4, parser.add_argument(
help="Weight for L2 loss") "--n-hidden", type=int, default=16, help="number of hidden gcn units"
parser.add_argument("--self-loop", action='store_true', )
help="graph self-loop (default=False)") parser.add_argument(
"--n-layers", type=int, default=1, help="number of hidden gcn layers"
)
parser.add_argument(
"--weight-decay", type=float, default=5e-4, help="Weight for L2 loss"
)
parser.add_argument(
"--self-loop",
action="store_true",
help="graph self-loop (default=False)",
)
parser.set_defaults(self_loop=False) parser.set_defaults(self_loop=False)
args = parser.parse_args() args = parser.parse_args()
print(args) print(args)
......
examples/tensorflow/rgcn/entity_classify.py
View file @ be8763fa
...@@ -9,16 +9,18 @@ Difference compared to tkipf/relation-gcn ...@@ -9,16 +9,18 @@ Difference compared to tkipf/relation-gcn
""" """
import argparse import argparse
import numpy as np
import time import time
from functools import partial
import numpy as np
import tensorflow as tf import tensorflow as tf
from model import BaseRGCN
from tensorflow.keras import layers from tensorflow.keras import layers
import dgl import dgl
from dgl.data.rdf import AIFBDataset, AMDataset, BGSDataset, MUTAGDataset
from dgl.nn.tensorflow import RelGraphConv from dgl.nn.tensorflow import RelGraphConv
from functools import partial
from dgl.data.rdf import AIFBDataset, MUTAGDataset, BGSDataset, AMDataset
from model import BaseRGCN
class EntityClassify(BaseRGCN): class EntityClassify(BaseRGCN):
def create_features(self): def create_features(self):
...@@ -26,19 +28,40 @@ class EntityClassify(BaseRGCN): ...@@ -26,19 +28,40 @@ class EntityClassify(BaseRGCN):
return features return features
def build_input_layer(self): def build_input_layer(self):
return RelGraphConv(self.num_nodes, self.h_dim, self.num_rels, "basis", return RelGraphConv(
self.num_bases, activation=tf.nn.relu, self_loop=self.use_self_loop, self.num_nodes,
dropout=self.dropout) self.h_dim,
self.num_rels,
"basis",
self.num_bases,
activation=tf.nn.relu,
self_loop=self.use_self_loop,
dropout=self.dropout,
)
def build_hidden_layer(self, idx): def build_hidden_layer(self, idx):
return RelGraphConv(self.h_dim, self.h_dim, self.num_rels, "basis", return RelGraphConv(
self.num_bases, activation=tf.nn.relu, self_loop=self.use_self_loop, self.h_dim,
dropout=self.dropout) self.h_dim,
self.num_rels,
"basis",
self.num_bases,
activation=tf.nn.relu,
self_loop=self.use_self_loop,
dropout=self.dropout,
)
def build_output_layer(self): def build_output_layer(self):
return RelGraphConv(self.h_dim, self.out_dim, self.num_rels, "basis", return RelGraphConv(
self.num_bases, activation=partial(tf.nn.softmax, axis=1), self.h_dim,
self_loop=self.use_self_loop) self.out_dim,
self.num_rels,
"basis",
self.num_bases,
activation=partial(tf.nn.softmax, axis=1),
self_loop=self.use_self_loop,
)
def acc(logits, labels, mask): def acc(logits, labels, mask):
logits = tf.gather(logits, mask) logits = tf.gather(logits, mask)
...@@ -47,15 +70,16 @@ def acc(logits, labels, mask): ...@@ -47,15 +70,16 @@ def acc(logits, labels, mask):
acc = tf.reduce_mean(tf.cast(indices == labels, dtype=tf.float32)) acc = tf.reduce_mean(tf.cast(indices == labels, dtype=tf.float32))
return acc return acc
def main(args): def main(args):
# load graph data # load graph data
if args.dataset == 'aifb': if args.dataset == "aifb":
dataset = AIFBDataset() dataset = AIFBDataset()
elif args.dataset == 'mutag': elif args.dataset == "mutag":
dataset = MUTAGDataset() dataset = MUTAGDataset()
elif args.dataset == 'bgs': elif args.dataset == "bgs":
dataset = BGSDataset() dataset = BGSDataset()
elif args.dataset == 'am': elif args.dataset == "am":
dataset = AMDataset() dataset = AMDataset()
else: else:
raise ValueError() raise ValueError()
...@@ -68,27 +92,27 @@ def main(args): ...@@ -68,27 +92,27 @@ def main(args):
num_rels = len(hg.canonical_etypes) num_rels = len(hg.canonical_etypes)
category = dataset.predict_category category = dataset.predict_category
num_classes = dataset.num_classes num_classes = dataset.num_classes
train_mask = hg.nodes[category].data.pop('train_mask') train_mask = hg.nodes[category].data.pop("train_mask")
test_mask = hg.nodes[category].data.pop('test_mask') test_mask = hg.nodes[category].data.pop("test_mask")
train_idx = tf.squeeze(tf.where(train_mask)) train_idx = tf.squeeze(tf.where(train_mask))
test_idx = tf.squeeze(tf.where(test_mask)) test_idx = tf.squeeze(tf.where(test_mask))
labels = hg.nodes[category].data.pop('labels') labels = hg.nodes[category].data.pop("labels")
# split dataset into train, validate, test # split dataset into train, validate, test
if args.validation: if args.validation:
val_idx = train_idx[:len(train_idx) // 5] val_idx = train_idx[: len(train_idx) // 5]
train_idx = train_idx[len(train_idx) // 5:] train_idx = train_idx[len(train_idx) // 5 :]
else: else:
val_idx = train_idx val_idx = train_idx
# calculate norm for each edge type and store in edge # calculate norm for each edge type and store in edge
for canonical_etype in hg.canonical_etypes: for canonical_etype in hg.canonical_etypes:
u, v, eid = hg.all_edges(form='all', etype=canonical_etype) u, v, eid = hg.all_edges(form="all", etype=canonical_etype)
_, inverse_index, count = tf.unique_with_counts(v) _, inverse_index, count = tf.unique_with_counts(v)
degrees = tf.gather(count, inverse_index) degrees = tf.gather(count, inverse_index)
norm = tf.ones(eid.shape[0]) / tf.cast(degrees, tf.float32) norm = tf.ones(eid.shape[0]) / tf.cast(degrees, tf.float32)
norm = tf.expand_dims(norm, 1) norm = tf.expand_dims(norm, 1)
hg.edges[canonical_etype].data['norm'] = norm hg.edges[canonical_etype].data["norm"] = norm
# get target category id # get target category id
category_id = len(hg.ntypes) category_id = len(hg.ntypes)
...@@ -97,7 +121,7 @@ def main(args): ...@@ -97,7 +121,7 @@ def main(args):
category_id = i category_id = i
# edge type and normalization factor # edge type and normalization factor
g = dgl.to_homogeneous(hg, edata=['norm']) g = dgl.to_homogeneous(hg, edata=["norm"])
# check cuda # check cuda
if args.gpu < 0: if args.gpu < 0:
...@@ -109,12 +133,12 @@ def main(args): ...@@ -109,12 +133,12 @@ def main(args):
use_cuda = True use_cuda = True
num_nodes = g.number_of_nodes() num_nodes = g.number_of_nodes()
node_ids = tf.range(num_nodes, dtype=tf.int64) node_ids = tf.range(num_nodes, dtype=tf.int64)
edge_norm = g.edata['norm'] edge_norm = g.edata["norm"]
edge_type = tf.cast(g.edata[dgl.ETYPE], tf.int64) edge_type = tf.cast(g.edata[dgl.ETYPE], tf.int64)
# find out the target node ids in g # find out the target node ids in g
node_tids = g.ndata[dgl.NTYPE] node_tids = g.ndata[dgl.NTYPE]
loc = (node_tids == category_id) loc = node_tids == category_id
target_idx = tf.squeeze(tf.where(loc)) target_idx = tf.squeeze(tf.where(loc))
# since the nodes are featureless, the input feature is then the node id. # since the nodes are featureless, the input feature is then the node id.
...@@ -122,38 +146,41 @@ def main(args): ...@@ -122,38 +146,41 @@ def main(args):
with tf.device(device): with tf.device(device):
# create model # create model
model = EntityClassify(num_nodes, model = EntityClassify(
args.n_hidden, num_nodes,
num_classes, args.n_hidden,
num_rels, num_classes,
num_bases=args.n_bases, num_rels,
num_hidden_layers=args.n_layers - 2, num_bases=args.n_bases,
dropout=args.dropout, num_hidden_layers=args.n_layers - 2,
use_self_loop=args.use_self_loop, dropout=args.dropout,
use_cuda=use_cuda) use_self_loop=args.use_self_loop,
use_cuda=use_cuda,
)
# optimizer # optimizer
optimizer = tf.keras.optimizers.Adam( optimizer = tf.keras.optimizers.Adam(learning_rate=args.lr)
learning_rate=args.lr)
# training loop # training loop
print("start training...") print("start training...")
forward_time = [] forward_time = []
backward_time = [] backward_time = []
loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy( loss_fcn = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=False) from_logits=False
)
for epoch in range(args.n_epochs): for epoch in range(args.n_epochs):
t0 = time.time() t0 = time.time()
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
logits = model(g, feats, edge_type, edge_norm) logits = model(g, feats, edge_type, edge_norm)
logits = tf.gather(logits, target_idx) logits = tf.gather(logits, target_idx)
loss = loss_fcn(tf.gather(labels, train_idx), tf.gather(logits, train_idx)) loss = loss_fcn(
tf.gather(labels, train_idx), tf.gather(logits, train_idx)
)
# Manually Weight Decay # Manually Weight Decay
# We found Tensorflow has a different implementation on weight decay # We found Tensorflow has a different implementation on weight decay
# of Adam(W) optimizer with PyTorch. And this results in worse results. # of Adam(W) optimizer with PyTorch. And this results in worse results.
# Manually adding weights to the loss to do weight decay solves this problem. # Manually adding weights to the loss to do weight decay solves this problem.
for weight in model.trainable_weights: for weight in model.trainable_weights:
loss = loss + \ loss = loss + args.l2norm * tf.nn.l2_loss(weight)
args.l2norm * tf.nn.l2_loss(weight)
t1 = time.time() t1 = time.time()
grads = tape.gradient(loss, model.trainable_weights) grads = tape.gradient(loss, model.trainable_weights)
optimizer.apply_gradients(zip(grads, model.trainable_weights)) optimizer.apply_gradients(zip(grads, model.trainable_weights))
...@@ -161,54 +188,93 @@ def main(args): ...@@ -161,54 +188,93 @@ def main(args):
forward_time.append(t1 - t0) forward_time.append(t1 - t0)
backward_time.append(t2 - t1) backward_time.append(t2 - t1)
print("Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}". print(
format(epoch, forward_time[-1], backward_time[-1])) "Epoch {:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}".format(
epoch, forward_time[-1], backward_time[-1]
)
)
train_acc = acc(logits, labels, train_idx) train_acc = acc(logits, labels, train_idx)
val_loss = loss_fcn(tf.gather(labels, val_idx), tf.gather(logits, val_idx)) val_loss = loss_fcn(
tf.gather(labels, val_idx), tf.gather(logits, val_idx)
)
val_acc = acc(logits, labels, val_idx) val_acc = acc(logits, labels, val_idx)
print("Train Accuracy: {:.4f} | Train Loss: {:.4f} | Validation Accuracy: {:.4f} | Validation loss: {:.4f}". print(
format(train_acc, loss.numpy().item(), val_acc, val_loss.numpy().item())) "Train Accuracy: {:.4f} | Train Loss: {:.4f} | Validation Accuracy: {:.4f} | Validation loss: {:.4f}".format(
train_acc,
loss.numpy().item(),
val_acc,
val_loss.numpy().item(),
)
)
print() print()
logits = model(g, feats, edge_type, edge_norm) logits = model(g, feats, edge_type, edge_norm)
logits = tf.gather(logits, target_idx) logits = tf.gather(logits, target_idx)
test_loss = loss_fcn(tf.gather(labels, test_idx), tf.gather(logits, test_idx)) test_loss = loss_fcn(
tf.gather(labels, test_idx), tf.gather(logits, test_idx)
)
test_acc = acc(logits, labels, test_idx) test_acc = acc(logits, labels, test_idx)
print("Test Accuracy: {:.4f} | Test loss: {:.4f}".format(test_acc, test_loss.numpy().item())) print(
"Test Accuracy: {:.4f} | Test loss: {:.4f}".format(
test_acc, test_loss.numpy().item()
)
)
print() print()
print("Mean forward time: {:4f}".format(np.mean(forward_time[len(forward_time) // 4:]))) print(
print("Mean backward time: {:4f}".format(np.mean(backward_time[len(backward_time) // 4:]))) "Mean forward time: {:4f}".format(
np.mean(forward_time[len(forward_time) // 4 :])
)
if __name__ == '__main__': )
parser = argparse.ArgumentParser(description='RGCN') print(
parser.add_argument("--dropout", type=float, default=0, "Mean backward time: {:4f}".format(
help="dropout probability") np.mean(backward_time[len(backward_time) // 4 :])
parser.add_argument("--n-hidden", type=int, default=16, )
help="number of hidden units") )
parser.add_argument("--gpu", type=int, default=-1,
help="gpu")
parser.add_argument("--lr", type=float, default=1e-2, if __name__ == "__main__":
help="learning rate") parser = argparse.ArgumentParser(description="RGCN")
parser.add_argument("--n-bases", type=int, default=-1, parser.add_argument(
help="number of filter weight matrices, default: -1 [use all]") "--dropout", type=float, default=0, help="dropout probability"
parser.add_argument("--n-layers", type=int, default=2, )
help="number of propagation rounds") parser.add_argument(
parser.add_argument("-e", "--n-epochs", type=int, default=50, "--n-hidden", type=int, default=16, help="number of hidden units"
help="number of training epochs") )
parser.add_argument("-d", "--dataset", type=str, required=True, parser.add_argument("--gpu", type=int, default=-1, help="gpu")
help="dataset to use") parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
parser.add_argument("--l2norm", type=float, default=0, parser.add_argument(
help="l2 norm coef") "--n-bases",
parser.add_argument("--use-self-loop", default=False, action='store_true', type=int,
help="include self feature as a special relation") default=-1,
help="number of filter weight matrices, default: -1 [use all]",
)
parser.add_argument(
"--n-layers", type=int, default=2, help="number of propagation rounds"
)
parser.add_argument(
"-e",
"--n-epochs",
type=int,
default=50,
help="number of training epochs",
)
parser.add_argument(
"-d", "--dataset", type=str, required=True, help="dataset to use"
)
parser.add_argument("--l2norm", type=float, default=0, help="l2 norm coef")
parser.add_argument(
"--use-self-loop",
default=False,
action="store_true",
help="include self feature as a special relation",
)
fp = parser.add_mutually_exclusive_group(required=False) fp = parser.add_mutually_exclusive_group(required=False)
fp.add_argument('--validation', dest='validation', action='store_true') fp.add_argument("--validation", dest="validation", action="store_true")
fp.add_argument('--testing', dest='validation', action='store_false') fp.add_argument("--testing", dest="validation", action="store_false")
parser.set_defaults(validation=True) parser.set_defaults(validation=True)
args = parser.parse_args() args = parser.parse_args()
print(args) print(args)
args.bfs_level = args.n_layers + 1 # pruning used nodes for memory args.bfs_level = args.n_layers + 1 # pruning used nodes for memory
main(args) main(args)
examples/tensorflow/rgcn/model.py
View file @ be8763fa
import tensorflow as tf import tensorflow as tf
from tensorflow.keras import layers from tensorflow.keras import layers
class BaseRGCN(layers.Layer): class BaseRGCN(layers.Layer):
def __init__(self, num_nodes, h_dim, out_dim, num_rels, num_bases, def __init__(
num_hidden_layers=1, dropout=0, self,
use_self_loop=False, use_cuda=False): num_nodes,
h_dim,
out_dim,
num_rels,
num_bases,
num_hidden_layers=1,
dropout=0,
use_self_loop=False,
use_cuda=False,
):
super(BaseRGCN, self).__init__() super(BaseRGCN, self).__init__()
self.num_nodes = num_nodes self.num_nodes = num_nodes
self.h_dim = h_dim self.h_dim = h_dim
......
examples/tensorflow/rgcn/utils.py
View file @ be8763fa
...@@ -7,6 +7,7 @@ https://github.com/MichSchli/RelationPrediction ...@@ -7,6 +7,7 @@ https://github.com/MichSchli/RelationPrediction
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
import dgl import dgl
####################################################################### #######################################################################
...@@ -15,11 +16,11 @@ import dgl ...@@ -15,11 +16,11 @@ import dgl
# #
####################################################################### #######################################################################
def get_adj_and_degrees(num_nodes, triplets): def get_adj_and_degrees(num_nodes, triplets):
""" Get adjacency list and degrees of the graph """Get adjacency list and degrees of the graph"""
"""
adj_list = [[] for _ in range(num_nodes)] adj_list = [[] for _ in range(num_nodes)]
for i,triplet in enumerate(triplets): for i, triplet in enumerate(triplets):
adj_list[triplet[0]].append([i, triplet[2]]) adj_list[triplet[0]].append([i, triplet[2]])
adj_list[triplet[2]].append([i, triplet[0]]) adj_list[triplet[2]].append([i, triplet[0]])
...@@ -27,6 +28,7 @@ def get_adj_and_degrees(num_nodes, triplets): ...@@ -27,6 +28,7 @@ def get_adj_and_degrees(num_nodes, triplets):
adj_list = [np.array(a) for a in adj_list] adj_list = [np.array(a) for a in adj_list]
return adj_list, degrees return adj_list, degrees
def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size): def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
"""Sample edges by neighborhool expansion. """Sample edges by neighborhool expansion.
...@@ -35,7 +37,7 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size): ...@@ -35,7 +37,7 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
""" """
edges = np.zeros((sample_size), dtype=np.int32) edges = np.zeros((sample_size), dtype=np.int32)
#initialize # initialize
sample_counts = np.array([d for d in degrees]) sample_counts = np.array([d for d in degrees])
picked = np.array([False for _ in range(n_triplets)]) picked = np.array([False for _ in range(n_triplets)])
seen = np.array([False for _ in degrees]) seen = np.array([False for _ in degrees])
...@@ -48,8 +50,9 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size): ...@@ -48,8 +50,9 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
weights[np.where(sample_counts == 0)] = 0 weights[np.where(sample_counts == 0)] = 0
probabilities = (weights) / np.sum(weights) probabilities = (weights) / np.sum(weights)
chosen_vertex = np.random.choice(np.arange(degrees.shape[0]), chosen_vertex = np.random.choice(
p=probabilities) np.arange(degrees.shape[0]), p=probabilities
)
chosen_adj_list = adj_list[chosen_vertex] chosen_adj_list = adj_list[chosen_vertex]
seen[chosen_vertex] = True seen[chosen_vertex] = True
...@@ -71,23 +74,36 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size): ...@@ -71,23 +74,36 @@ def sample_edge_neighborhood(adj_list, degrees, n_triplets, sample_size):
return edges return edges
def sample_edge_uniform(adj_list, degrees, n_triplets, sample_size): def sample_edge_uniform(adj_list, degrees, n_triplets, sample_size):
"""Sample edges uniformly from all the edges.""" """Sample edges uniformly from all the edges."""
all_edges = np.arange(n_triplets) all_edges = np.arange(n_triplets)
return np.random.choice(all_edges, sample_size, replace=False) return np.random.choice(all_edges, sample_size, replace=False)
def generate_sampled_graph_and_labels(triplets, sample_size, split_size,
num_rels, adj_list, degrees, def generate_sampled_graph_and_labels(
negative_rate, sampler="uniform"): triplets,
sample_size,
split_size,
num_rels,
adj_list,
degrees,
negative_rate,
sampler="uniform",
):
"""Get training graph and signals """Get training graph and signals
First perform edge neighborhood sampling on graph, then perform negative First perform edge neighborhood sampling on graph, then perform negative
sampling to generate negative samples sampling to generate negative samples
""" """
# perform edge neighbor sampling # perform edge neighbor sampling
if sampler == "uniform": if sampler == "uniform":
edges = sample_edge_uniform(adj_list, degrees, len(triplets), sample_size) edges = sample_edge_uniform(
adj_list, degrees, len(triplets), sample_size
)
elif sampler == "neighbor": elif sampler == "neighbor":
edges = sample_edge_neighborhood(adj_list, degrees, len(triplets), sample_size) edges = sample_edge_neighborhood(
adj_list, degrees, len(triplets), sample_size
)
else: else:
raise ValueError("Sampler type must be either 'uniform' or 'neighbor'.") raise ValueError("Sampler type must be either 'uniform' or 'neighbor'.")
...@@ -99,14 +115,16 @@ def generate_sampled_graph_and_labels(triplets, sample_size, split_size, ...@@ -99,14 +115,16 @@ def generate_sampled_graph_and_labels(triplets, sample_size, split_size,
relabeled_edges = np.stack((src, rel, dst)).transpose() relabeled_edges = np.stack((src, rel, dst)).transpose()
# negative sampling # negative sampling
samples, labels = negative_sampling(relabeled_edges, len(uniq_v), samples, labels = negative_sampling(
negative_rate) relabeled_edges, len(uniq_v), negative_rate
)
# further split graph, only half of the edges will be used as graph # further split graph, only half of the edges will be used as graph
# structure, while the rest half is used as unseen positive samples # structure, while the rest half is used as unseen positive samples
split_size = int(sample_size * split_size) split_size = int(sample_size * split_size)
graph_split_ids = np.random.choice(np.arange(sample_size), graph_split_ids = np.random.choice(
size=split_size, replace=False) np.arange(sample_size), size=split_size, replace=False
)
src = src[graph_split_ids] src = src[graph_split_ids]
dst = dst[graph_split_ids] dst = dst[graph_split_ids]
rel = rel[graph_split_ids] rel = rel[graph_split_ids]
...@@ -114,10 +132,12 @@ def generate_sampled_graph_and_labels(triplets, sample_size, split_size, ...@@ -114,10 +132,12 @@ def generate_sampled_graph_and_labels(triplets, sample_size, split_size,
# build DGL graph # build DGL graph
print("# sampled nodes: {}".format(len(uniq_v))) print("# sampled nodes: {}".format(len(uniq_v)))
print("# sampled edges: {}".format(len(src) * 2)) print("# sampled edges: {}".format(len(src) * 2))
g, rel, norm = build_graph_from_triplets(len(uniq_v), num_rels, g, rel, norm = build_graph_from_triplets(
(src, rel, dst)) len(uniq_v), num_rels, (src, rel, dst)
)
return g, uniq_v, rel, norm, samples, labels return g, uniq_v, rel, norm, samples, labels
def comp_deg_norm(g): def comp_deg_norm(g):
g = g.local_var() g = g.local_var()
in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy() in_deg = g.in_degrees(range(g.number_of_nodes())).float().numpy()
...@@ -125,11 +145,12 @@ def comp_deg_norm(g): ...@@ -125,11 +145,12 @@ def comp_deg_norm(g):
norm[np.isinf(norm)] = 0 norm[np.isinf(norm)] = 0
return norm return norm
def build_graph_from_triplets(num_nodes, num_rels, triplets): def build_graph_from_triplets(num_nodes, num_rels, triplets):
""" Create a DGL graph. The graph is bidirectional because RGCN authors """Create a DGL graph. The graph is bidirectional because RGCN authors
use reversed relations. use reversed relations.
This function also generates edge type and normalization factor This function also generates edge type and normalization factor
(reciprocal of node incoming degree) (reciprocal of node incoming degree)
""" """
g = dgl.DGLGraph() g = dgl.DGLGraph()
g.add_nodes(num_nodes) g.add_nodes(num_nodes)
...@@ -143,17 +164,19 @@ def build_graph_from_triplets(num_nodes, num_rels, triplets): ...@@ -143,17 +164,19 @@ def build_graph_from_triplets(num_nodes, num_rels, triplets):
print("# nodes: {}, # edges: {}".format(num_nodes, len(src))) print("# nodes: {}, # edges: {}".format(num_nodes, len(src)))
return g, rel, norm return g, rel, norm
def build_test_graph(num_nodes, num_rels, edges): def build_test_graph(num_nodes, num_rels, edges):
src, rel, dst = edges.transpose() src, rel, dst = edges.transpose()
print("Test graph:") print("Test graph:")
return build_graph_from_triplets(num_nodes, num_rels, (src, rel, dst)) return build_graph_from_triplets(num_nodes, num_rels, (src, rel, dst))
def negative_sampling(pos_samples, num_entity, negative_rate): def negative_sampling(pos_samples, num_entity, negative_rate):
size_of_batch = len(pos_samples) size_of_batch = len(pos_samples)
num_to_generate = size_of_batch * negative_rate num_to_generate = size_of_batch * negative_rate
neg_samples = np.tile(pos_samples, (negative_rate, 1)) neg_samples = np.tile(pos_samples, (negative_rate, 1))
labels = np.zeros(size_of_batch * (negative_rate + 1), dtype=np.float32) labels = np.zeros(size_of_batch * (negative_rate + 1), dtype=np.float32)
labels[: size_of_batch] = 1 labels[:size_of_batch] = 1
values = np.random.randint(num_entity, size=num_to_generate) values = np.random.randint(num_entity, size=num_to_generate)
choices = np.random.uniform(size=num_to_generate) choices = np.random.uniform(size=num_to_generate)
subj = choices > 0.5 subj = choices > 0.5
...@@ -162,4 +185,3 @@ def negative_sampling(pos_samples, num_entity, negative_rate): ...@@ -162,4 +185,3 @@ def negative_sampling(pos_samples, num_entity, negative_rate):
neg_samples[obj, 2] = values[obj] neg_samples[obj, 2] = values[obj]
return np.concatenate((pos_samples, neg_samples)), labels return np.concatenate((pos_samples, neg_samples)), labels
examples/tensorflow/sgc/sgc.py
View file @ be8763fa
...@@ -16,11 +16,10 @@ import tensorflow_addons as tfa ...@@ -16,11 +16,10 @@ import tensorflow_addons as tfa
from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
from dgl.nn.tensorflow.conv import SGConv from dgl.nn.tensorflow.conv import SGConv
_DATASETS = { _DATASETS = {
'citeseer': CiteseerGraphDataset(verbose=False), "citeseer": CiteseerGraphDataset(verbose=False),
'cora': CoraGraphDataset(verbose=False), "cora": CoraGraphDataset(verbose=False),
'pubmed': PubmedGraphDataset(verbose=False) "pubmed": PubmedGraphDataset(verbose=False),
} }
...@@ -29,7 +28,7 @@ def load_data(dataset): ...@@ -29,7 +28,7 @@ def load_data(dataset):
def _sum_boolean_tensor(x): def _sum_boolean_tensor(x):
return tf.reduce_sum(tf.cast(x, dtype='int64')) return tf.reduce_sum(tf.cast(x, dtype="int64"))
def describe_data(data): def describe_data(data):
...@@ -38,9 +37,9 @@ def describe_data(data): ...@@ -38,9 +37,9 @@ def describe_data(data):
n_edges = g.number_of_edges() n_edges = g.number_of_edges()
num_classes = data.num_classes num_classes = data.num_classes
train_mask = g.ndata['train_mask'] train_mask = g.ndata["train_mask"]
val_mask = g.ndata['val_mask'] val_mask = g.ndata["val_mask"]
test_mask = g.ndata['test_mask'] test_mask = g.ndata["test_mask"]
description = textwrap.dedent( description = textwrap.dedent(
f""" f"""
...@@ -65,7 +64,7 @@ class SGC(tf.keras.Model): ...@@ -65,7 +64,7 @@ class SGC(tf.keras.Model):
out_feats=self.num_classes, out_feats=self.num_classes,
k=2, k=2,
cached=True, cached=True,
bias=bias bias=bias,
) )
def call(self, inputs): def call(self, inputs):
...@@ -73,7 +72,7 @@ class SGC(tf.keras.Model): ...@@ -73,7 +72,7 @@ class SGC(tf.keras.Model):
@property @property
def in_feats(self): def in_feats(self):
return self.g.ndata['feat'].shape[1] return self.g.ndata["feat"].shape[1]
@property @property
def num_nodes(self): def num_nodes(self):
...@@ -87,7 +86,7 @@ class SGC(tf.keras.Model): ...@@ -87,7 +86,7 @@ class SGC(tf.keras.Model):
def train_step(self, data): def train_step(self, data):
X, y = data X, y = data
mask = self.g.ndata['train_mask'] mask = self.g.ndata["train_mask"]
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
y_pred = self(X, training=True) y_pred = self(X, training=True)
...@@ -101,7 +100,7 @@ class SGC(tf.keras.Model): ...@@ -101,7 +100,7 @@ class SGC(tf.keras.Model):
def test_step(self, data): def test_step(self, data):
X, y = data X, y = data
mask = self.g.ndata['val_mask'] mask = self.g.ndata["val_mask"]
y_pred = self(X, training=False) y_pred = self(X, training=False)
self.compiled_loss(y[mask], y_pred[mask]) self.compiled_loss(y[mask], y_pred[mask])
self.compiled_metrics.update_state(y[mask], y_pred[mask]) self.compiled_metrics.update_state(y[mask], y_pred[mask])
...@@ -111,12 +110,12 @@ class SGC(tf.keras.Model): ...@@ -111,12 +110,12 @@ class SGC(tf.keras.Model):
super().compile(*args, **kwargs, run_eagerly=True) super().compile(*args, **kwargs, run_eagerly=True)
def fit(self, *args, **kwargs): def fit(self, *args, **kwargs):
kwargs['batch_size'] = self.num_nodes kwargs["batch_size"] = self.num_nodes
kwargs['shuffle'] = False kwargs["shuffle"] = False
super().fit(*args, **kwargs) super().fit(*args, **kwargs)
def predict(self, *args, **kwargs): def predict(self, *args, **kwargs):
kwargs['batch_size'] = self.num_nodes kwargs["batch_size"] = self.num_nodes
return super().predict(*args, **kwargs) return super().predict(*args, **kwargs)
...@@ -125,41 +124,32 @@ def main(dataset, lr, bias, n_epochs, weight_decay): ...@@ -125,41 +124,32 @@ def main(dataset, lr, bias, n_epochs, weight_decay):
print(describe_data(data)) print(describe_data(data))
g = data[0] g = data[0]
X = g.ndata['feat'] X = g.ndata["feat"]
y = g.ndata['label'] y = g.ndata["label"]
model = SGC(g=g, num_classes=data.num_classes, bias=bias) model = SGC(g=g, num_classes=data.num_classes, bias=bias)
loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True) loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tfa.optimizers.AdamW(weight_decay, lr) optimizer = tfa.optimizers.AdamW(weight_decay, lr)
accuracy = tf.metrics.SparseCategoricalAccuracy(name='accuracy') accuracy = tf.metrics.SparseCategoricalAccuracy(name="accuracy")
model.compile(optimizer, loss, metrics=[accuracy]) model.compile(optimizer, loss, metrics=[accuracy])
model.fit( model.fit(x=X, y=y, epochs=n_epochs, validation_data=(X, y))
x=X,
y=y,
epochs=n_epochs,
validation_data=(X, y)
)
y_pred = model.predict(X, batch_size=len(X)) y_pred = model.predict(X, batch_size=len(X))
test_mask = g.ndata['test_mask'] test_mask = g.ndata["test_mask"]
test_accuracy = accuracy(y[test_mask], y_pred[test_mask]) test_accuracy = accuracy(y[test_mask], y_pred[test_mask])
print(f"Test Accuracy: {test_accuracy:.1%}") print(f"Test Accuracy: {test_accuracy:.1%}")
def _parse_args(): def _parse_args():
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description='Run experiment for Simple Graph Convolution (SGC)' description="Run experiment for Simple Graph Convolution (SGC)"
)
parser.add_argument(
"--dataset", default='cora', help="dataset to run"
)
parser.add_argument(
"--lr", type=float, default=0.2, help="learning rate"
) )
parser.add_argument("--dataset", default="cora", help="dataset to run")
parser.add_argument("--lr", type=float, default=0.2, help="learning rate")
parser.add_argument( parser.add_argument(
"--bias", action='store_true', default=False, help="flag to use bias" "--bias", action="store_true", default=False, help="flag to use bias"
) )
parser.add_argument( parser.add_argument(
"--n-epochs", type=int, default=100, help="number of training epochs" "--n-epochs", type=int, default=100, help="number of training epochs"
...@@ -170,12 +160,12 @@ def _parse_args(): ...@@ -170,12 +160,12 @@ def _parse_args():
return parser.parse_args() return parser.parse_args()
if __name__ == '__main__': if __name__ == "__main__":
args = _parse_args() args = _parse_args()
main( main(
dataset=args.dataset, dataset=args.dataset,
lr=args.lr, lr=args.lr,
bias=args.bias, bias=args.bias,
n_epochs=args.n_epochs, n_epochs=args.n_epochs,
weight_decay=args.weight_decay weight_decay=args.weight_decay,
) )
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment