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