Unverified Commit 4ec8f204 authored by Mufei Li's avatar Mufei Li Committed by GitHub
Browse files

Update (#1319)

parent 066d290f
...@@ -6,7 +6,7 @@ Graph Isomorphism Network (GIN) ...@@ -6,7 +6,7 @@ Graph Isomorphism Network (GIN)
Dependencies Dependencies
------------ ------------
- PyTorch 1.0.1+ - PyTorch 1.1.0+
- sklearn - sklearn
- tqdm - tqdm
......
...@@ -19,7 +19,7 @@ def collate(samples): ...@@ -19,7 +19,7 @@ def collate(samples):
for g in graphs: for g in graphs:
# deal with node feats # deal with node feats
for key in g.node_attr_schemes().keys(): for key in g.node_attr_schemes().keys():
g.ndata[key] = torch.from_numpy(g.ndata[key]).float() g.ndata[key] = g.ndata[key].float()
# no edge feats # no edge feats
batched_graph = dgl.batch(graphs) batched_graph = dgl.batch(graphs)
labels = torch.tensor(labels) labels = torch.tensor(labels)
......
...@@ -73,14 +73,14 @@ def eval_net(args, net, dataloader, criterion): ...@@ -73,14 +73,14 @@ def eval_net(args, net, dataloader, criterion):
def main(args): def main(args):
# set up seeds, args.seed supported # set up seeds, args.seed supported
torch.manual_seed(seed=0) torch.manual_seed(seed=args.seed)
np.random.seed(seed=0) np.random.seed(seed=args.seed)
is_cuda = not args.disable_cuda and torch.cuda.is_available() is_cuda = not args.disable_cuda and torch.cuda.is_available()
if is_cuda: if is_cuda:
args.device = torch.device("cuda:" + str(args.device)) args.device = torch.device("cuda:" + str(args.device))
torch.cuda.manual_seed_all(seed=0) torch.cuda.manual_seed_all(seed=args.seed)
else: else:
args.device = torch.device("cpu") args.device = torch.device("cpu")
...@@ -109,9 +109,9 @@ def main(args): ...@@ -109,9 +109,9 @@ def main(args):
lrbar = tqdm(range(args.epochs), unit="epoch", position=5, ncols=0, file=sys.stdout) lrbar = tqdm(range(args.epochs), unit="epoch", position=5, ncols=0, file=sys.stdout)
for epoch, _, _ in zip(tbar, vbar, lrbar): for epoch, _, _ in zip(tbar, vbar, lrbar):
scheduler.step()
train(args, model, trainloader, optimizer, criterion, epoch) train(args, model, trainloader, optimizer, criterion, epoch)
scheduler.step()
train_loss, train_acc = eval_net( train_loss, train_acc = eval_net(
args, model, trainloader, criterion) args, model, trainloader, criterion)
......
...@@ -19,6 +19,7 @@ class Parser(): ...@@ -19,6 +19,7 @@ class Parser():
# dataset # dataset
self.parser.add_argument( self.parser.add_argument(
'--dataset', type=str, default="MUTAG", '--dataset', type=str, default="MUTAG",
choices=['MUTAG', 'COLLAB', 'IMDBBINARY', 'IMDBMULTI'],
help='name of dataset (default: MUTAG)') help='name of dataset (default: MUTAG)')
self.parser.add_argument( self.parser.add_argument(
'--batch_size', type=int, default=32, '--batch_size', type=int, default=32,
...@@ -39,9 +40,6 @@ class Parser(): ...@@ -39,9 +40,6 @@ class Parser():
help='which gpu device to use (default: 0)') help='which gpu device to use (default: 0)')
# net # net
self.parser.add_argument(
'--net', type=str, default="gin",
help='gnn net (default: gin)')
self.parser.add_argument( self.parser.add_argument(
'--num_layers', type=int, default=5, '--num_layers', type=int, default=5,
help='number of layers (default: 5)') help='number of layers (default: 5)')
...@@ -64,9 +62,6 @@ class Parser(): ...@@ -64,9 +62,6 @@ class Parser():
self.parser.add_argument( self.parser.add_argument(
'--learn_eps', action="store_true", '--learn_eps', action="store_true",
help='learn the epsilon weighting') help='learn the epsilon weighting')
self.parser.add_argument(
'--degree_as_tag', action="store_true",
help='take the degree of nodes as input feature')
# learning # learning
self.parser.add_argument( self.parser.add_argument(
......
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