def main(args): # load graph data if args.dataset == 'aifb': dataset = AIFB() elif args.dataset == 'mutag': dataset = MUTAG() elif args.dataset == 'bgs': dataset = BGS() elif args.dataset == 'am': dataset = AM() else: raise ValueError() g = dataset.graph category = dataset.predict_category num_classes = dataset.num_classes test_idx = dataset.test_idx labels = dataset.labels category_id = len(g.ntypes) for i, ntype in enumerate(g.ntypes): if ntype == category: category_id = i # check cuda use_cuda = args.gpu >= 0 and th.cuda.is_available() if use_cuda: th.cuda.set_device(args.gpu) labels = labels.cuda() test_idx = test_idx.cuda() # create model model = EntityClassify(g, args.n_hidden, num_classes, num_bases=args.n_bases, num_hidden_layers=args.n_layers - 2, use_self_loop=args.use_self_loop) # training loop model.load_state_dict(th.load(args.model_path)) if use_cuda: model.cuda() print("start testing...") model.eval() logits = model.forward()[category_id] test_loss = F.cross_entropy(logits[test_idx], labels[test_idx]) test_acc = th.sum(logits[test_idx].argmax( dim=1) == labels[test_idx]).item() / len(test_idx) print("Test Acc: {:.4f} | Test loss: {:.4f}".format( test_acc, test_loss.item())) print()
def main(args): # load graph data if args.dataset == 'aifb': dataset = AIFB() elif args.dataset == 'mutag': dataset = MUTAG() elif args.dataset == 'bgs': dataset = BGS() elif args.dataset == 'am': dataset = AM() else: raise ValueError() g = dataset.graph category = dataset.predict_category num_classes = dataset.num_classes train_idx = dataset.train_idx test_idx = dataset.test_idx labels = dataset.labels category_id = len(g.ntypes) for i, ntype in enumerate(g.ntypes): if ntype == category: category_id = i # 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:] else: val_idx = train_idx # check cuda use_cuda = args.gpu >= 0 and th.cuda.is_available() if use_cuda: th.cuda.set_device(args.gpu) labels = labels.cuda() train_idx = train_idx.cuda() test_idx = test_idx.cuda() # create model model = EntityClassify(g, args.n_hidden, num_classes, num_bases=args.n_bases, num_hidden_layers=args.n_layers - 2, dropout=args.dropout, use_self_loop=args.use_self_loop) if use_cuda: model.cuda() # optimizer optimizer = th.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2norm) # training loop print("start training...") dur = [] model.train() for epoch in range(args.n_epochs): optimizer.zero_grad() if epoch > 5: t0 = time.time() logits = model()[category] loss = F.cross_entropy(logits[train_idx], labels[train_idx]) loss.backward() optimizer.step() t1 = time.time() if epoch > 5: dur.append(t1 - t0) train_acc = th.sum(logits[train_idx].argmax(dim=1) == labels[train_idx]).item() / len(train_idx) val_loss = F.cross_entropy(logits[val_idx], labels[val_idx]) val_acc = th.sum(logits[val_idx].argmax(dim=1) == labels[val_idx]).item() / len(val_idx) print("Epoch {:05d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}". format(epoch, train_acc, loss.item(), val_acc, val_loss.item(), np.average(dur))) print() if args.model_path is not None: th.save(model.state_dict(), args.model_path) model.eval() logits = model.forward()[category] test_loss = F.cross_entropy(logits[test_idx], labels[test_idx]) test_acc = th.sum(logits[test_idx].argmax(dim=1) == labels[test_idx]).item() / len(test_idx) print("Test Acc: {:.4f} | Test loss: {:.4f}".format(test_acc, test_loss.item())) print()
def main(args): # load graph data if args.dataset == 'aifb': dataset = AIFB() elif args.dataset == 'mutag': dataset = MUTAG() elif args.dataset == 'bgs': dataset = BGS() elif args.dataset == 'am': dataset = AM() else: raise ValueError() g = dataset.graph category = dataset.predict_category num_classes = dataset.num_classes train_idx = dataset.train_idx test_idx = dataset.test_idx labels = dataset.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:] else: val_idx = train_idx # check cuda use_cuda = args.gpu >= 0 and th.cuda.is_available() if use_cuda: th.cuda.set_device(args.gpu) train_label = labels[train_idx] val_label = labels[val_idx] test_label = labels[test_idx] # create embeddings embed_layer = RelGraphEmbed(g, args.n_hidden) node_embed = embed_layer() # create model model = EntityClassify(g, args.n_hidden, num_classes, num_bases=args.n_bases, num_hidden_layers=args.n_layers - 2, dropout=args.dropout, use_self_loop=args.use_self_loop) if use_cuda: model.cuda() # train sampler sampler = HeteroNeighborSampler(g, category, [args.fanout] * args.n_layers) loader = DataLoader(dataset=train_idx.numpy(), batch_size=args.batch_size, collate_fn=sampler.sample_blocks, shuffle=True, num_workers=0) # validation sampler val_sampler = HeteroNeighborSampler(g, category, [None] * args.n_layers) _, val_blocks = val_sampler.sample_blocks(val_idx) # test sampler test_sampler = HeteroNeighborSampler(g, category, [None] * args.n_layers) _, test_blocks = test_sampler.sample_blocks(test_idx) # optimizer all_params = itertools.chain(model.parameters(), embed_layer.parameters()) optimizer = th.optim.Adam(all_params, lr=args.lr, weight_decay=args.l2norm) # training loop print("start training...") dur = [] for epoch in range(args.n_epochs): model.train() optimizer.zero_grad() if epoch > 3: t0 = time.time() for i, (seeds, blocks) in enumerate(loader): batch_tic = time.time() emb = extract_embed(node_embed, blocks[0]) lbl = labels[seeds[category]] if use_cuda: emb = {k: e.cuda() for k, e in emb.items()} lbl = lbl.cuda() logits = model(emb, blocks)[category] loss = F.cross_entropy(logits, lbl) loss.backward() optimizer.step() train_acc = th.sum(logits.argmax(dim=1) == lbl).item() / len( seeds[category]) print( "Epoch {:05d} | Batch {:03d} | Train Acc: {:.4f} | Train Loss: {:.4f} | Time: {:.4f}" .format(epoch, i, train_acc, loss.item(), time.time() - batch_tic)) if epoch > 3: dur.append(time.time() - t0) val_loss, val_acc = evaluate(model, val_idx, val_blocks, node_embed, labels, category, use_cuda) print( "Epoch {:05d} | Valid Acc: {:.4f} | Valid loss: {:.4f} | Time: {:.4f}" .format(epoch, val_acc, val_loss.item(), np.average(dur))) print() if args.model_path is not None: th.save(model.state_dict(), args.model_path) test_loss, test_acc = evaluate(model, test_idx, test_blocks, node_embed, labels, category, use_cuda) print("Test Acc: {:.4f} | Test loss: {:.4f}".format( test_acc, test_loss.item())) print()
def main(args, devices): # load graph data ogb_dataset = False if args.dataset == 'aifb': dataset = AIFB() elif args.dataset == 'mutag': dataset = MUTAG() elif args.dataset == 'bgs': dataset = BGS() elif args.dataset == 'am': dataset = AM() else: raise ValueError() # Load from hetero-graph hg = dataset.graph num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) category = dataset.predict_category num_classes = dataset.num_classes train_idx = dataset.train_idx test_idx = dataset.test_idx labels = dataset.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:] else: val_idx = train_idx # calculate norm for each edge type and store in edge for canonical_etypes in hg.canonical_etypes: u, v, eid = hg.all_edges(form='all', etype=canonical_etypes) _, inverse_index, count = th.unique(v, return_inverse=True, return_counts=True) degrees = count[inverse_index] norm = th.ones(eid.shape[0]) / degrees norm = norm.unsqueeze(1) hg.edges[canonical_etypes].data['norm'] = norm # get target category id category_id = len(hg.ntypes) for i, ntype in enumerate(hg.ntypes): if ntype == category: category_id = i g = dgl.to_homo(hg) g.ndata[dgl.NTYPE].share_memory_() g.edata[dgl.ETYPE].share_memory_() g.edata['norm'].share_memory_() node_ids = th.arange(g.number_of_nodes()) # find out the target node ids node_tids = g.ndata[dgl.NTYPE] loc = (node_tids == category_id) target_idx = node_ids[loc] target_idx.share_memory_() n_gpus = len(devices) # cpu if devices[0] == -1: run(0, 0, args, ['cpu'], (g, num_of_ntype, num_classes, num_rels, target_idx, train_idx, val_idx, test_idx, labels)) # gpu elif n_gpus == 1: run(0, n_gpus, args, devices, (g, num_of_ntype, num_classes, num_rels, target_idx, train_idx, val_idx, test_idx, labels)) # multi gpu else: procs = [] num_train_seeds = train_idx.shape[0] tseeds_per_proc = num_train_seeds // n_gpus for proc_id in range(n_gpus): proc_train_seeds = train_idx[proc_id * tseeds_per_proc : (proc_id + 1) * tseeds_per_proc \ if (proc_id + 1) * tseeds_per_proc < num_train_seeds \ else num_train_seeds] p = mp.Process(target=run, args=(proc_id, n_gpus, args, devices, (g, num_of_ntype, num_classes, num_rels, target_idx, proc_train_seeds, val_idx, test_idx, labels))) p.start() procs.append(p) for p in procs: p.join()