Skip to content

GitLab

Unverified Commit 6124667f authored by Da Zheng's avatar Da Zheng Committed by GitHub
Browse files

[Model] Refactor the sampling examples (#498) 

* reorganize sampling code.

* speedup gcn_ns.

* speed up gcn_cv

* fix graphsage_cv.

* undo the modification.

* accel training.

* update readme.
parent a3febc06
Hide whitespace changes
Inline Side-by-side
Showing with 166 additions and 276 deletions
examples/mxnet/sampling/README.md
View file @ 6124667f
......@@ -15,44 +15,44 @@ pip install mxnet --pre
### Neighbor Sampling & Skip Connection
cora: test accuracy ~83% with `--num-neighbors 2`, ~84% by training on the full graph
```
DGLBACKEND=mxnet python gcn_ns_sc.py --dataset cora --self-loop --num-neighbors 2 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_ns --dataset cora --self-loop --num-neighbors 2 --batch-size 1000 --test-batch-size 5000
```
citeseer: test accuracy ~69% with `--num-neighbors 2`, ~70% by training on the full graph
```
DGLBACKEND=mxnet python gcn_ns_sc.py --dataset citeseer --self-loop --num-neighbors 2 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_ns --dataset citeseer --self-loop --num-neighbors 2 --batch-size 1000 --test-batch-size 5000
```
pubmed: test accuracy ~76% with `--num-neighbors 3`, ~77% by training on the full graph
pubmed: test accuracy ~78% with `--num-neighbors 3`, ~77% by training on the full graph
```
DGLBACKEND=mxnet python gcn_ns_sc.py --dataset pubmed --self-loop --num-neighbors 3 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_ns --dataset pubmed --self-loop --num-neighbors 3 --batch-size 1000 --test-batch-size 5000
```
reddit: test accuracy ~91% with `--num-neighbors 2` and `--batch-size 1000`, ~93% by training on the full graph
reddit: test accuracy ~91% with `--num-neighbors 3` and `--batch-size 1000`, ~93% by training on the full graph
```
DGLBACKEND=mxnet python gcn_ns_sc.py --dataset reddit-self-loop --num-neighbors 2 --batch-size 1000 --test-batch-size 500 --n-hidden 64
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_ns --dataset reddit-self-loop --num-neighbors 3 --batch-size 1000 --test-batch-size 5000 --n-hidden 64
```
### Control Variate & Skip Connection
cora: test accuracy ~84% with `--num-neighbors 1`, ~84% by training on the full graph
```
DGLBACKEND=mxnet python gcn_cv_sc.py --dataset cora --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_cv --dataset cora --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000
```
citeseer: test accuracy ~69% with `--num-neighbors 1`, ~70% by training on the full graph
```
DGLBACKEND=mxnet python gcn_cv_sc.py --dataset citeseer --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_cv --dataset citeseer --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000
```
pubmed: test accuracy ~77% with `--num-neighbors 1`, ~77% by training on the full graph
pubmed: test accuracy ~79% with `--num-neighbors 1`, ~77% by training on the full graph
```
DGLBACKEND=mxnet python gcn_cv_sc.py --dataset pubmed --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000 --gpu 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_cv --dataset pubmed --self-loop --num-neighbors 1 --batch-size 1000000 --test-batch-size 1000000
```
reddit: test accuracy ~93% with `--num-neighbors 1` and `--batch-size 1000`, ~93% by training on the full graph
```
DGLBACKEND=mxnet python gcn_cv_sc.py --dataset reddit-self-loop --num-neighbors 1 --batch-size 1000 --test-batch-size 500 --n-hidden 64
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model gcn_cv --dataset reddit-self-loop --num-neighbors 1 --batch-size 10000 --test-batch-size 5000 --n-hidden 64
```
### Control Variate & GraphSAGE-mean
......@@ -61,6 +61,6 @@ Following [Control Variate](https://arxiv.org/abs/1710.10568), we use the mean p
reddit: test accuracy 96.1% with `--num-neighbors 1` and `--batch-size 1000`, ~96.2% in [Control Variate](https://arxiv.org/abs/1710.10568) with `--num-neighbors 2` and `--batch-size 1000`
```
DGLBACKEND=mxnet python graphsage_cv.py --batch-size 1000 --test-batch-size 500 --n-epochs 50 --dataset reddit --num-neighbors 1 --n-hidden 128 --dropout 0.2 --weight-decay 0
DGLBACKEND=mxnet python examples/mxnet/sampling/train.py --model graphsage_cv --batch-size 1000 --test-batch-size 5000 --n-epochs 50 --dataset reddit --num-neighbors 1 --n-hidden 128 --dropout 0.2 --weight-decay 0
```
examples/mxnet/sampling/gcn_cv_sc.py
View file @ 6124667f
......@@ -21,14 +21,15 @@ class NodeUpdate(gluon.Block):
def forward(self, node):
h = node.data['h']
norm = node.data['norm']
if self.test:
norm = node.data['norm']
h = h * norm
else:
agg_history_str = 'agg_h_{}'.format(self.layer_id-1)
agg_history = node.data[agg_history_str]
subg_norm = node.data['subg_norm']
# control variate
h = h + agg_history
h = h * subg_norm + agg_history * norm
if self.dropout:
h = mx.nd.Dropout(h, p=self.dropout)
h = self.dense(h)
......@@ -39,7 +40,6 @@ class NodeUpdate(gluon.Block):
return {'activation': h}
class GCNSampling(gluon.Block):
def __init__(self,
in_feats,
......@@ -85,7 +85,7 @@ class GCNSampling(gluon.Block):
nf.layers[i].data['h'] = h
nf.block_compute(i,
fn.copy_src(src='h', out='m'),
lambda node : {'h': node.mailbox['m'].mean(axis=1)},
fn.sum(msg='m', out='h'),
layer)
h = nf.layers[i+1].data.pop('activation')
# update history
......@@ -139,63 +139,24 @@ class GCNInfer(gluon.Block):
return h
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.gpu >= 0:
ctx = mx.gpu(args.gpu)
else:
ctx = mx.cpu()
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
train_nid = mx.nd.array(np.nonzero(data.train_mask)[0]).astype(np.int64)
test_nid = mx.nd.array(np.nonzero(data.test_mask)[0]).astype(np.int64)
num_neighbors = args.num_neighbors
n_layers = args.n_layers
features = mx.nd.array(data.features).as_in_context(ctx)
labels = mx.nd.array(data.labels).as_in_context(ctx)
train_mask = mx.nd.array(data.train_mask).as_in_context(ctx)
val_mask = mx.nd.array(data.val_mask).as_in_context(ctx)
test_mask = mx.nd.array(data.test_mask).as_in_context(ctx)
def gcn_cv_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples):
features = g.ndata['features']
labels = g.ndata['labels']
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().asscalar()
n_test_samples = test_mask.sum().asscalar()
n_val_samples = val_mask.sum().asscalar()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
n_train_samples,
n_val_samples,
n_test_samples))
# create GCN model
g = DGLGraph(data.graph, readonly=True)
g.ndata['features'] = features
norm = mx.nd.expand_dims(1./g.in_degrees().astype('float32'), 1)
g.ndata['norm'] = norm.as_in_context(ctx)
degs = g.in_degrees().astype('float32').asnumpy()
degs[degs > args.num_neighbors] = args.num_neighbors
g.ndata['subg_norm'] = mx.nd.expand_dims(mx.nd.array(1./degs, ctx=ctx), 1)
g.update_all(fn.copy_src(src='features', out='m'),
fn.sum(msg='m', out='preprocess'),
lambda node : {'preprocess': node.data['preprocess'] * node.data['norm']})
n_layers = args.n_layers
for i in range(n_layers):
g.ndata['h_{}'.format(i)] = mx.nd.zeros((features.shape[0], args.n_hidden), ctx=ctx)
g.ndata['h_{}'.format(n_layers-1)] = mx.nd.zeros((features.shape[0], 2*args.n_hidden), ctx=ctx)
model = GCNSampling(in_feats,
......@@ -229,21 +190,21 @@ def main(args):
dur = []
for epoch in range(args.n_epochs):
for nf in dgl.contrib.sampling.NeighborSampler(g, args.batch_size,
num_neighbors,
args.num_neighbors,
neighbor_type='in',
num_workers=32,
shuffle=True,
num_hops=n_layers,
seed_nodes=train_nid):
for i in range(n_layers):
agg_history_str = 'agg_h_{}'.format(i)
g.pull(nf.layer_parent_nid(i+1), fn.copy_src(src='h_{}'.format(i), out='m'),
fn.sum(msg='m', out=agg_history_str),
lambda node : {agg_history_str: node.data[agg_history_str] * node.data['norm']})
fn.sum(msg='m', out=agg_history_str))
node_embed_names = [['preprocess', 'h_0']]
for i in range(1, n_layers):
node_embed_names.append(['h_{}'.format(i), 'agg_h_{}'.format(i-1)])
node_embed_names.append(['agg_h_{}'.format(n_layers-1)])
node_embed_names.append(['h_{}'.format(i), 'agg_h_{}'.format(i-1), 'subg_norm', 'norm'])
node_embed_names.append(['agg_h_{}'.format(n_layers-1), 'subg_norm', 'norm'])
nf.copy_from_parent(node_embed_names=node_embed_names)
# forward
......@@ -269,6 +230,7 @@ def main(args):
trainer._kvstore.pull(idx, out=infer_params[key].data())
num_acc = 0.
num_tests = 0
for nf in dgl.contrib.sampling.NeighborSampler(g, args.test_batch_size,
g.number_of_nodes(),
......@@ -284,38 +246,9 @@ def main(args):
batch_nids = nf.layer_parent_nid(-1).as_in_context(ctx)
batch_labels = labels[batch_nids]
num_acc += (pred.argmax(axis=1) == batch_labels).sum().asscalar()
num_tests += nf.layer_size(-1)
break
print("Test Accuracy {:.4f}". format(num_acc/num_tests))
print("Test Accuracy {:.4f}". format(num_acc/n_test_samples))
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=3e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--batch-size", type=int, default=1000,
help="train batch size")
parser.add_argument("--test-batch-size", type=int, default=1000,
help="test batch size")
parser.add_argument("--num-neighbors", type=int, default=2,
help="number of neighbors to be sampled")
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("--self-loop", action='store_true',
help="graph self-loop (default=False)")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
args = parser.parse_args()
print(args)
main(args)
examples/mxnet/sampling/gcn_ns_sc.py
View file @ 6124667f
......@@ -10,17 +10,15 @@ from dgl.data import register_data_args, load_data
class NodeUpdate(gluon.Block):
def __init__(self, in_feats, out_feats, activation=None, test=False, concat=False):
def __init__(self, in_feats, out_feats, activation=None, concat=False):
super(NodeUpdate, self).__init__()
self.dense = gluon.nn.Dense(out_feats, in_units=in_feats)
self.activation = activation
self.concat = concat
self.test = test
def forward(self, node):
h = node.data['h']
if self.test:
h = h * node.data['norm']
h = h * node.data['norm']
h = self.dense(h)
# skip connection
if self.concat:
......@@ -63,9 +61,11 @@ class GCNSampling(gluon.Block):
if self.dropout:
h = mx.nd.Dropout(h, p=self.dropout)
nf.layers[i].data['h'] = h
degs = nf.layer_in_degree(i + 1).astype('float32').as_in_context(h.context)
nf.layers[i + 1].data['norm'] = mx.nd.expand_dims(1./degs, 1)
nf.block_compute(i,
fn.copy_src(src='h', out='m'),
lambda node : {'h': node.mailbox['m'].mean(axis=1)},
fn.sum(msg='m', out='h'),
layer)
h = nf.layers[-1].data.pop('activation')
......@@ -86,13 +86,13 @@ class GCNInfer(gluon.Block):
self.layers = gluon.nn.Sequential()
# input layer
skip_start = (0 == n_layers-1)
self.layers.add(NodeUpdate(in_feats, n_hidden, activation, test=True, concat=skip_start))
self.layers.add(NodeUpdate(in_feats, n_hidden, activation, concat=skip_start))
# hidden layers
for i in range(1, n_layers):
skip_start = (i == n_layers-1)
self.layers.add(NodeUpdate(n_hidden, n_hidden, activation, test=True, concat=skip_start))
self.layers.add(NodeUpdate(n_hidden, n_hidden, activation, concat=skip_start))
# output layer
self.layers.add(NodeUpdate(2*n_hidden, n_classes, test=True))
self.layers.add(NodeUpdate(2*n_hidden, n_classes))
def forward(self, nf):
......@@ -106,55 +106,12 @@ class GCNInfer(gluon.Block):
fn.sum(msg='m', out='h'),
layer)
h = nf.layers[-1].data.pop('activation')
return h
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.gpu >= 0:
ctx = mx.gpu(args.gpu)
else:
ctx = mx.cpu()
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
return nf.layers[-1].data.pop('activation')
train_nid = mx.nd.array(np.nonzero(data.train_mask)[0]).astype(np.int64).as_in_context(ctx)
test_nid = mx.nd.array(np.nonzero(data.test_mask)[0]).astype(np.int64).as_in_context(ctx)
features = mx.nd.array(data.features).as_in_context(ctx)
labels = mx.nd.array(data.labels).as_in_context(ctx)
train_mask = mx.nd.array(data.train_mask).as_in_context(ctx)
val_mask = mx.nd.array(data.val_mask).as_in_context(ctx)
test_mask = mx.nd.array(data.test_mask).as_in_context(ctx)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().asscalar()
n_val_samples = val_mask.sum().asscalar()
n_test_samples = test_mask.sum().asscalar()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
n_train_samples,
n_val_samples,
n_test_samples))
# create GCN model
g = DGLGraph(data.graph, readonly=True)
g.ndata['features'] = features
num_neighbors = args.num_neighbors
def gcn_ns_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples):
in_feats = g.ndata['features'].shape[1]
labels = g.ndata['labels']
degs = g.in_degrees().astype('float32').as_in_context(ctx)
norm = mx.nd.expand_dims(1./degs, 1)
......@@ -189,15 +146,13 @@ def main(args):
# initialize graph
dur = []
for epoch in range(args.n_epochs):
index = 0
for nf in dgl.contrib.sampling.NeighborSampler(g, args.batch_size,
args.num_neighbors,
neighbor_type='in',
shuffle=True,
num_workers=32,
num_hops=args.n_layers+1,
seed_nodes=train_nid):
print(index)
index = index + 1
nf.copy_from_parent()
# forward
with mx.autograd.record():
......@@ -217,6 +172,7 @@ def main(args):
trainer._kvstore.pull(idx, out=infer_params[key].data())
num_acc = 0.
num_tests = 0
for nf in dgl.contrib.sampling.NeighborSampler(g, args.test_batch_size,
g.number_of_nodes(),
......@@ -228,38 +184,7 @@ def main(args):
batch_nids = nf.layer_parent_nid(-1).astype('int64').as_in_context(ctx)
batch_labels = labels[batch_nids]
num_acc += (pred.argmax(axis=1) == batch_labels).sum().asscalar()
num_tests += nf.layer_size(-1)
break
print("Test Accuracy {:.4f}". format(num_acc/n_test_samples))
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=3e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--batch-size", type=int, default=1000,
help="batch size")
parser.add_argument("--test-batch-size", type=int, default=1000,
help="test batch size")
parser.add_argument("--num-neighbors", type=int, default=3,
help="number of neighbors to be sampled")
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("--self-loop", action='store_true',
help="graph self-loop (default=False)")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
args = parser.parse_args()
print(args)
main(args)
print("Test Accuracy {:.4f}". format(num_acc/num_tests))
examples/mxnet/sampling/graphsage_cv.py
View file @ 6124667f
......@@ -114,7 +114,7 @@ class GraphSAGETrain(gluon.Block):
for i, layer in enumerate(self.layers):
parent_nid = dgl.utils.toindex(nf.layer_parent_nid(i+1))
layer_nid = nf.map_from_parent_nid(i, parent_nid)
layer_nid = nf.map_from_parent_nid(i, parent_nid).as_in_context(h.context)
self_h = h[layer_nid]
# activation from previous layer of myself, used in graphSAGE
nf.layers[i+1].data['self_h'] = self_h
......@@ -168,7 +168,7 @@ class GraphSAGEInfer(gluon.Block):
for i, layer in enumerate(self.layers):
nf.layers[i].data['h'] = h
parent_nid = dgl.utils.toindex(nf.layer_parent_nid(i+1))
layer_nid = nf.map_from_parent_nid(i, parent_nid)
layer_nid = nf.map_from_parent_nid(i, parent_nid).as_in_context(h.context)
# activation from previous layer of the nodes in (i+1)-th layer, used in graphSAGE
self_h = h[layer_nid]
nf.layers[i+1].data['self_h'] = self_h
......@@ -181,64 +181,23 @@ class GraphSAGEInfer(gluon.Block):
return h
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.gpu >= 0:
ctx = mx.gpu(args.gpu)
else:
ctx = mx.cpu()
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
train_nid = mx.nd.array(np.nonzero(data.train_mask)[0]).astype(np.int64)
test_nid = mx.nd.array(np.nonzero(data.test_mask)[0]).astype(np.int64)
num_neighbors = args.num_neighbors
n_layers = args.n_layers
features = mx.nd.array(data.features).as_in_context(ctx)
labels = mx.nd.array(data.labels).as_in_context(ctx)
train_mask = mx.nd.array(data.train_mask).as_in_context(ctx)
val_mask = mx.nd.array(data.val_mask).as_in_context(ctx)
test_mask = mx.nd.array(data.test_mask).as_in_context(ctx)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().asscalar()
n_test_samples = test_mask.sum().asscalar()
n_val_samples = val_mask.sum().asscalar()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
n_train_samples,
n_val_samples,
n_test_samples))
g = DGLGraph(data.graph, readonly=True)
g.ndata['features'] = features
def graphsage_cv_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples):
features = g.ndata['features']
labels = g.ndata['labels']
in_feats = g.ndata['features'].shape[1]
norm = mx.nd.expand_dims(1./g.in_degrees().astype('float32'), 1)
g.ndata['norm'] = norm.as_in_context(ctx)
degs = g.in_degrees().astype('float32').asnumpy()
degs[degs > num_neighbors] = num_neighbors
degs[degs > args.num_neighbors] = args.num_neighbors
g.ndata['subg_norm'] = mx.nd.expand_dims(mx.nd.array(1./degs, ctx=ctx), 1)
g.update_all(fn.copy_src(src='features', out='m'),
fn.sum(msg='m', out='preprocess'),
lambda node : {'preprocess': node.data['preprocess'] * node.data['norm']})
n_layers = args.n_layers
for i in range(n_layers):
g.ndata['h_{}'.format(i)] = mx.nd.zeros((features.shape[0], args.n_hidden), ctx=ctx)
......@@ -271,9 +230,10 @@ def main(args):
dur = []
for epoch in range(args.n_epochs):
for nf in dgl.contrib.sampling.NeighborSampler(g, args.batch_size,
num_neighbors,
args.num_neighbors,
neighbor_type='in',
shuffle=True,
num_workers=32,
num_hops=n_layers,
add_self_loop=True,
seed_nodes=train_nid):
......@@ -311,6 +271,7 @@ def main(args):
trainer._kvstore.pull(idx, out=infer_params[key].data())
num_acc = 0.
num_tests = 0
for nf in dgl.contrib.sampling.NeighborSampler(g, args.test_batch_size,
g.number_of_nodes(),
......@@ -327,38 +288,7 @@ def main(args):
batch_nids = nf.layer_parent_nid(-1).as_in_context(ctx)
batch_labels = labels[batch_nids]
num_acc += (pred.argmax(axis=1) == batch_labels).sum().asscalar()
num_tests += nf.layer_size(-1)
break
print("Test Accuracy {:.4f}". format(num_acc/n_test_samples))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GraphSAGE with Control Variate')
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=3e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--batch-size", type=int, default=1000,
help="train batch size")
parser.add_argument("--test-batch-size", type=int, default=1000,
help="test batch size")
parser.add_argument("--num-neighbors", type=int, default=3,
help="number of neighbors to be sampled")
parser.add_argument("--n-hidden", type=int, default=16,
help="number of hidden GraphSAGE units")
parser.add_argument("--n-layers", type=int, default=1,
help="number of hidden GraphSAGE layers")
parser.add_argument("--self-loop", action='store_true',
help="graph self-loop (default=False)")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
args = parser.parse_args()
print(args)
main(args)
print("Test Accuracy {:.4f}". format(num_acc/num_tests))
examples/mxnet/sampling/train.py 0 → 100644
View file @ 6124667f
import argparse, time, math
import numpy as np
import mxnet as mx
from mxnet import gluon
from functools import partial
import dgl
import dgl.function as fn
from dgl import DGLGraph
from dgl.data import register_data_args, load_data
from gcn_ns_sc import gcn_ns_train
from gcn_cv_sc import gcn_cv_train
from graphsage_cv import graphsage_cv_train
def main(args):
# load and preprocess dataset
data = load_data(args)
if args.gpu >= 0:
ctx = mx.gpu(args.gpu)
else:
ctx = mx.cpu()
if args.self_loop and not args.dataset.startswith('reddit'):
data.graph.add_edges_from([(i,i) for i in range(len(data.graph))])
train_nid = mx.nd.array(np.nonzero(data.train_mask)[0]).astype(np.int64).as_in_context(ctx)
test_nid = mx.nd.array(np.nonzero(data.test_mask)[0]).astype(np.int64).as_in_context(ctx)
features = mx.nd.array(data.features).as_in_context(ctx)
labels = mx.nd.array(data.labels).as_in_context(ctx)
train_mask = mx.nd.array(data.train_mask).as_in_context(ctx)
val_mask = mx.nd.array(data.val_mask).as_in_context(ctx)
test_mask = mx.nd.array(data.test_mask).as_in_context(ctx)
in_feats = features.shape[1]
n_classes = data.num_labels
n_edges = data.graph.number_of_edges()
n_train_samples = train_mask.sum().asscalar()
n_val_samples = val_mask.sum().asscalar()
n_test_samples = test_mask.sum().asscalar()
print("""----Data statistics------'
#Edges %d
#Classes %d
#Train samples %d
#Val samples %d
#Test samples %d""" %
(n_edges, n_classes,
n_train_samples,
n_val_samples,
n_test_samples))
# create GCN model
g = DGLGraph(data.graph, readonly=True)
g.ndata['features'] = features
g.ndata['labels'] = labels
if args.model == "gcn_ns":
gcn_ns_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples)
elif args.model == "gcn_cv":
gcn_cv_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples)
elif args.model == "graphsage_cv":
graphsage_cv_train(g, ctx, args, n_classes, train_nid, test_nid, n_test_samples)
else:
print("unknown model. Please choose from gcn_ns, gcn_cv, graphsage_cv")
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='GCN')
register_data_args(parser)
parser.add_argument("--model", type=str,
help="select a model. Valid models: gcn_ns, gcn_cv, graphsage_cv")
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=3e-2,
help="learning rate")
parser.add_argument("--n-epochs", type=int, default=200,
help="number of training epochs")
parser.add_argument("--batch-size", type=int, default=1000,
help="batch size")
parser.add_argument("--test-batch-size", type=int, default=1000,
help="test batch size")
parser.add_argument("--num-neighbors", type=int, default=3,
help="number of neighbors to be sampled")
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("--self-loop", action='store_true',
help="graph self-loop (default=False)")
parser.add_argument("--weight-decay", type=float, default=5e-4,
help="Weight for L2 loss")
parser.add_argument("--nworkers", type=int, default=1,
help="number of workers")
args = parser.parse_args()
print(args)
main(args)
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