def forward(self, g, x=None): """Forward pass. Parameters ---------- g : `dgl.DGLHeteroGraph`, input graph Returns ------- g : `dgl.DGLHeteroGraph` output graph """ import dgl # get homogeneous subgraph g_ = dgl.to_homo(g.edge_type_subgraph(["n1_neighbors_n1"])) if x is None: # get node attributes x = g.nodes["n1"].data["h0"] x = self.f_in(x) # message passing on h**o graph x = self._sequential(g_, x) # put attribute back in the graph g.nodes["n1"].data["h"] = x return g
def construct_homo_from_hetero_dglgraph(g): homo_g = dgl.to_homo(g) # still a heterograph in dgl hg = dgl.DGLGraph() hg.add_nodes(homo_g.number_of_nodes('_N')) u, v = homo_g.all_edges(form='uv', order='eid') hg.add_edges(u, v, homo_g.edata) hg.readonly(True) # sampler needs readonly graph return hg
def main(args): # load graph data if args.dataset == 'aifb': dataset = AIFBDataset() elif args.dataset == 'mutag': dataset = MUTAGDataset() elif args.dataset == 'bgs': dataset = BGSDataset() elif args.dataset == 'am': dataset = AMDataset() else: raise ValueError() # Load from hetero-graph hg = dataset[0] num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) 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_idx = mx.nd.array(np.nonzero(train_mask.asnumpy())[0], dtype='int64') test_idx = mx.nd.array(np.nonzero(test_mask.asnumpy())[0], dtype='int64') labels = mx.nd.array(hg.nodes[category].data.pop('labels'), dtype='int64') # 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_etype in hg.canonical_etypes: u, v, eid = hg.all_edges(form='all', etype=canonical_etype) v = v.asnumpy() _, inverse_index, count = np.unique(v, return_inverse=True, return_counts=True) degrees = count[inverse_index] norm = np.ones(eid.shape[0]) / degrees hg.edges[canonical_etype].data['norm'] = mx.nd.expand_dims( mx.nd.array(norm), axis=1) # 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) num_nodes = g.number_of_nodes() node_ids = mx.nd.arange(num_nodes) edge_norm = g.edata['norm'] edge_type = g.edata[dgl.ETYPE] # find out the target node ids in g node_tids = g.ndata[dgl.NTYPE] loc = (node_tids == category_id) loc = mx.nd.array(np.nonzero(loc.asnumpy())[0], dtype='int64') target_idx = node_ids[loc] # since the nodes are featureless, the input feature is then the node id. feats = mx.nd.arange(num_nodes, dtype='int32') # check cuda use_cuda = args.gpu >= 0 if use_cuda: ctx = mx.gpu(args.gpu) feats = feats.as_in_context(ctx) edge_type = edge_type.as_in_context(ctx) edge_norm = edge_norm.as_in_context(ctx) labels = labels.as_in_context(ctx) train_idx = train_idx.as_in_context(ctx) g = g.to(ctx) else: ctx = mx.cpu(0) # 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, gpu_id=args.gpu) model.initialize(ctx=ctx) # optimizer trainer = gluon.Trainer(model.collect_params(), 'adam', { 'learning_rate': args.lr, 'wd': args.l2norm }) loss_fcn = gluon.loss.SoftmaxCELoss(from_logits=False) # training loop print("start training...") forward_time = [] backward_time = [] for epoch in range(args.n_epochs): t0 = time.time() with mx.autograd.record(): pred = model(g, feats, edge_type, edge_norm) pred = pred[target_idx] loss = loss_fcn(pred[train_idx], labels[train_idx]) t1 = time.time() loss.backward() trainer.step(len(train_idx)) t2 = time.time() 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])) train_acc = F.sum( mx.nd.cast(pred[train_idx].argmax(axis=1), 'int64') == labels[train_idx]).asscalar() / train_idx.shape[0] val_acc = F.sum( mx.nd.cast(pred[val_idx].argmax( axis=1), 'int64') == labels[val_idx]).asscalar() / len(val_idx) print("Train Accuracy: {:.4f} | Validation Accuracy: {:.4f}".format( train_acc, val_acc)) print() logits = model.forward(g, feats, edge_type, edge_norm) logits = logits[target_idx] test_acc = F.sum( mx.nd.cast(logits[test_idx].argmax( axis=1), 'int64') == labels[test_idx]).asscalar() / len(test_idx) print("Test Accuracy: {:.4f}".format(test_acc)) 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:])))
def main(args, devices): # load graph data ogb_dataset = False if args.dataset == 'aifb': dataset = AIFBDataset() elif args.dataset == 'mutag': dataset = MUTAGDataset() elif args.dataset == 'bgs': dataset = BGSDataset() elif args.dataset == 'am': dataset = AMDataset() else: raise ValueError() # Load from hetero-graph hg = dataset[0] num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) 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') labels = hg.nodes[category].data.pop('labels') train_idx = th.nonzero(train_mask).squeeze() test_idx = th.nonzero(test_mask).squeeze() # 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_etype in hg.canonical_etypes: u, v, eid = hg.all_edges(form='all', etype=canonical_etype) _, 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_etype].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()
def test_convert(): hg = create_test_heterograph() hs = [] for ntype in hg.ntypes: h = F.randn((hg.number_of_nodes(ntype), 5)) hg.nodes[ntype].data['h'] = h hs.append(h) hg.nodes['user'].data['x'] = F.randn((3, 3)) ws = [] for etype in hg.canonical_etypes: w = F.randn((hg.number_of_edges(etype), 5)) hg.edges[etype].data['w'] = w ws.append(w) hg.edges['plays'].data['x'] = F.randn((4, 3)) g = dgl.to_homo(hg) assert F.array_equal(F.cat(hs, dim=0), g.ndata['h']) assert 'x' not in g.ndata assert F.array_equal(F.cat(ws, dim=0), g.edata['w']) assert 'x' not in g.edata src, dst = g.all_edges(order='eid') src = F.asnumpy(src) dst = F.asnumpy(dst) etype_id, eid = F.asnumpy(g.edata[dgl.ETYPE]), F.asnumpy(g.edata[dgl.EID]) ntype_id, nid = F.asnumpy(g.ndata[dgl.NTYPE]), F.asnumpy(g.ndata[dgl.NID]) for i in range(g.number_of_edges()): srctype = hg.ntypes[ntype_id[src[i]]] dsttype = hg.ntypes[ntype_id[dst[i]]] etype = hg.etypes[etype_id[i]] src_i, dst_i = hg.find_edges([eid[i]], (srctype, etype, dsttype)) assert np.asscalar(F.asnumpy(src_i)) == nid[src[i]] assert np.asscalar(F.asnumpy(dst_i)) == nid[dst[i]] mg = nx.MultiDiGraph([('user', 'user', 'follows'), ('user', 'game', 'plays'), ('user', 'game', 'wishes'), ('developer', 'game', 'develops')]) for _mg in [None, mg]: hg2 = dgl.to_hetero(g, ['user', 'game', 'developer'], ['follows', 'plays', 'wishes', 'develops'], ntype_field=dgl.NTYPE, etype_field=dgl.ETYPE, metagraph=_mg) assert set(hg.ntypes) == set(hg2.ntypes) assert set(hg.canonical_etypes) == set(hg2.canonical_etypes) for ntype in hg.ntypes: assert hg.number_of_nodes(ntype) == hg2.number_of_nodes(ntype) assert F.array_equal(hg.nodes[ntype].data['h'], hg2.nodes[ntype].data['h']) for canonical_etype in hg.canonical_etypes: src, dst = hg.all_edges(etype=canonical_etype, order='eid') src2, dst2 = hg2.all_edges(etype=canonical_etype, order='eid') assert F.array_equal(src, src2) assert F.array_equal(dst, dst2) assert F.array_equal(hg.edges[canonical_etype].data['w'], hg2.edges[canonical_etype].data['w']) # hetero_from_homo test case 2 g = dgl.graph([(0, 2), (1, 2), (2, 3), (0, 3)]) g.ndata[dgl.NTYPE] = F.tensor([0, 0, 1, 2]) g.edata[dgl.ETYPE] = F.tensor([0, 0, 1, 2]) hg = dgl.to_hetero(g, ['l0', 'l1', 'l2'], ['e0', 'e1', 'e2']) assert set(hg.canonical_etypes) == set([('l0', 'e0', 'l1'), ('l1', 'e1', 'l2'), ('l0', 'e2', 'l2')]) assert hg.number_of_nodes('l0') == 2 assert hg.number_of_nodes('l1') == 1 assert hg.number_of_nodes('l2') == 1 assert hg.number_of_edges('e0') == 2 assert hg.number_of_edges('e1') == 1 assert hg.number_of_edges('e2') == 1 # hetero_from_homo test case 3 mg = nx.MultiDiGraph([('user', 'movie', 'watches'), ('user', 'TV', 'watches')]) g = dgl.graph([(0, 1), (0, 2)]) g.ndata[dgl.NTYPE] = F.tensor([0, 1, 2]) g.edata[dgl.ETYPE] = F.tensor([0, 0]) for _mg in [None, mg]: hg = dgl.to_hetero(g, ['user', 'TV', 'movie'], ['watches'], metagraph=_mg) assert set(hg.canonical_etypes) == set([('user', 'watches', 'movie'), ('user', 'watches', 'TV')]) assert hg.number_of_nodes('user') == 1 assert hg.number_of_nodes('TV') == 1 assert hg.number_of_nodes('movie') == 1 assert hg.number_of_edges(('user', 'watches', 'TV')) == 1 assert hg.number_of_edges(('user', 'watches', 'movie')) == 1 assert len(hg.etypes) == 2 # hetero_to_homo test case 2 hg = dgl.bipartite([(0, 0), (1, 1)], card=(2, 3)) g = dgl.to_homo(hg) assert g.number_of_nodes() == 5
def main(args, devices): # load graph data ogb_dataset = False if args.dataset == 'aifb': dataset = AIFBDataset() elif args.dataset == 'mutag': dataset = MUTAGDataset() elif args.dataset == 'bgs': dataset = BGSDataset() elif args.dataset == 'am': dataset = AMDataset() elif args.dataset == 'ogbn-mag': dataset = DglNodePropPredDataset(name=args.dataset) ogb_dataset = True else: raise ValueError() if ogb_dataset is True: split_idx = dataset.get_idx_split() train_idx = split_idx["train"]['paper'] val_idx = split_idx["valid"]['paper'] test_idx = split_idx["test"]['paper'] hg_orig, labels = dataset[0] subgs = {} for etype in hg_orig.canonical_etypes: u, v = hg_orig.all_edges(etype=etype) subgs[etype] = (u, v) subgs[(etype[2], 'rev-' + etype[1], etype[0])] = (v, u) hg = dgl.heterograph(subgs) hg.nodes['paper'].data['feat'] = hg_orig.nodes['paper'].data['feat'] labels = labels['paper'].squeeze() num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) num_classes = dataset.num_classes if args.dataset == 'ogbn-mag': category = 'paper' print('Number of relations: {}'.format(num_rels)) print('Number of class: {}'.format(num_classes)) print('Number of train: {}'.format(len(train_idx))) print('Number of valid: {}'.format(len(val_idx))) print('Number of test: {}'.format(len(test_idx))) if args.node_feats: node_feats = [] for ntype in hg.ntypes: if len(hg.nodes[ntype].data) == 0: node_feats.append(None) else: assert len(hg.nodes[ntype].data) == 1 feat = hg.nodes[ntype].data.pop('feat') node_feats.append(feat.share_memory_()) else: node_feats = [None] * num_of_ntype else: # Load from hetero-graph hg = dataset[0] num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) 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') labels = hg.nodes[category].data.pop('labels') train_idx = th.nonzero(train_mask).squeeze() test_idx = th.nonzero(test_mask).squeeze() node_feats = [None] * num_of_ntype # AIFB, MUTAG, BGS and AM datasets do not provide validation set split. # Split train set into train and validation if args.validation is set # otherwise use train set as the validation set. 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 if args.global_norm is False: for canonical_etype in hg.canonical_etypes: u, v, eid = hg.all_edges(form='all', etype=canonical_etype) _, 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_etype].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) if args.global_norm: u, v, eid = g.all_edges(form='all') _, 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) g.edata['norm'] = norm 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_() train_idx.share_memory_() val_idx.share_memory_() test_idx.share_memory_() n_gpus = len(devices) # cpu if devices[0] == -1: run(0, 0, args, ['cpu'], (g, node_feats, num_of_ntype, num_classes, num_rels, target_idx, train_idx, val_idx, test_idx, labels), None, None) # gpu elif n_gpus == 1: run(0, n_gpus, args, devices, (g, node_feats, num_of_ntype, num_classes, num_rels, target_idx, train_idx, val_idx, test_idx, labels), None, None) # multi gpu else: queue = mp.Queue(n_gpus) procs = [] num_train_seeds = train_idx.shape[0] num_valid_seeds = val_idx.shape[0] num_test_seeds = test_idx.shape[0] train_seeds = th.randperm(num_train_seeds) valid_seeds = th.randperm(num_valid_seeds) test_seeds = th.randperm(num_test_seeds) tseeds_per_proc = num_train_seeds // n_gpus vseeds_per_proc = num_valid_seeds // n_gpus tstseeds_per_proc = num_test_seeds // n_gpus for proc_id in range(n_gpus): # we have multi-gpu for training, evaluation and testing # so split trian set, valid set and test set into num-of-gpu parts. proc_train_seeds = train_seeds[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] proc_valid_seeds = valid_seeds[proc_id * vseeds_per_proc : (proc_id + 1) * vseeds_per_proc \ if (proc_id + 1) * vseeds_per_proc < num_valid_seeds \ else num_valid_seeds] proc_test_seeds = test_seeds[proc_id * tstseeds_per_proc : (proc_id + 1) * tstseeds_per_proc \ if (proc_id + 1) * tstseeds_per_proc < num_test_seeds \ else num_test_seeds] p = mp.Process(target=run, args=(proc_id, n_gpus, args, devices, (g, node_feats, num_of_ntype, num_classes, num_rels, target_idx, train_idx, val_idx, test_idx, labels), (proc_train_seeds, proc_valid_seeds, proc_test_seeds), queue)) p.start() procs.append(p) for p in procs: p.join()
import dgl if __name__ == "__main__": follows_g = dgl.graph([(0, 1), (1, 2)], 'user', 'follows') devs_g = dgl.bipartite([(0, 0), (1, 1)], 'developer', 'develops', 'game') hetero_g = dgl.hetero_from_relations([follows_g, devs_g]) homo_g = dgl.to_homo(hetero_g) hetero_g_2 = dgl.to_hetero(homo_g, hetero_g.ntypes, hetero_g.etypes) print(hetero_g) print(hetero_g_2) print("here")
def construct_dglgraph(data_path, file_prefix='test', emb_path='./emb/', load_version='', h**o=False, datainput='narr argrole temp'): labels = json.load( open(os.path.join(data_path, file_prefix + '.label.json'), 'r')) event_file = os.path.join( emb_path, load_version + '.' + file_prefix + '.eventemb.npy') if os.access(event_file, os.F_OK): entity_file = os.path.join( emb_path, load_version + '.' + file_prefix + '.entityemb.npy') else: event_file = os.path.join(data_path, file_prefix + '.eventemb.npy') entity_file = os.path.join(data_path, file_prefix + '.entityemb.npy') event_emb = np.load(event_file) entity_emb = np.load(entity_file) print('Graph Construction: loaded %s and %s' % (event_file, entity_file)) # construct hetero graph dgld = dict() dgld[('event', 'event2entity', 'entity')] = [(t[0], t[1]) for t in labels['event2entity']] # dgld[('event', 'event2event', 'event')] = [(t[0], t[1]) for t in labels['event2event']] dgld[('event', 'narr', 'event')] = [(t[0], t[1]) for t in labels['narr']] dgld[('event', 'coref', 'event')] = [(t[0], t[1]) for t in labels['coref']] # dgld[('event', 'eer', 'event')] = [(t[0], t[1]) for t in labels['eer']] dgld[('event', 'argrole', 'entity')] = [(t[0], t[1]) for t in labels['argrole']] dgld[('event', 'temp', 'event')] = [(t[0], t[1]) for t in labels['temp']] dgld[('entity', 'entity2entity', 'entity')] = [(t[0], t[1]) for t in labels['entity2entity']] g = dgl.heterograph(dgld) g.nodes['event'].data['x'] = torch.tensor(event_emb) g.nodes['entity'].data['x'] = torch.tensor(entity_emb) g.nodes['event'].data['y'] = torch.tensor([t[0] for t in labels['event']]) g.nodes['entity'].data['y'] = torch.tensor( [t[0] for t in labels['entity']]) g.edges['event2entity'].data['y'] = torch.tensor( [t[2] for t in labels['event2entity']]) g.edges['entity2entity'].data['y'] = torch.tensor( [t[2] for t in labels['entity2entity']]) #### add event/entity id #### event_cnt = event_emb.shape[0] entity_cnt = entity_emb.shape[0] g.nodes['event'].data['id'] = torch.tensor(list(range(event_cnt))) g.nodes['entity'].data['id'] = torch.tensor( [i + event_cnt for i in range(entity_cnt)]) # event-event relations cotrol by using string variable if 'narr' in datainput: g.edges['narr'].data['y'] = torch.tensor( [t[2] for t in labels['narr']]) if 'coref' in datainput: g.edges['coref'].data['y'] = torch.tensor( [t[2] for t in labels['coref']]) if 'eer' in datainput: g.edges['eer'].data['y'] = torch.tensor([t[2] for t in labels['eer']]) if 'argrole' in datainput: g.edges['argrole'].data['y'] = torch.tensor( [t[2] for t in labels['argrole']]) if 'temp' in datainput: g.edges['temp'].data['y'] = torch.tensor( [t[2] for t in labels['temp']]) if h**o: g.nodes['entity'].data['y'] += type_class_dict['event'] homo_g = dgl.to_homo(g) g = dgl.DGLGraph() g.add_nodes(homo_g.number_of_nodes('_N'), homo_g.ndata) u, v = homo_g.all_edges(form='uv', order='eid') g.add_edges(u, v, homo_g.edata) g.readonly(True) #sampler needs readonly graph return g
def main(args): # load graph data if args.dataset == 'aifb': dataset = AIFBDataset() elif args.dataset == 'mutag': dataset = MUTAGDataset() elif args.dataset == 'bgs': dataset = BGSDataset() elif args.dataset == 'am': dataset = AMDataset() else: raise ValueError() # preprocessing in cpu with tf.device("/cpu:0"): # Load from hetero-graph hg = dataset[0] num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) 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_idx = tf.squeeze(tf.where(train_mask)) test_idx = tf.squeeze(tf.where(test_mask)) 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:] 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) _, 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 # get target category id category_id = len(hg.ntypes) for i, ntype in enumerate(hg.ntypes): if ntype == category: category_id = i # edge type and normalization factor g = dgl.to_homo(hg) # check cuda if args.gpu < 0: device = "/cpu:0" use_cuda = False else: device = "/gpu:{}".format(args.gpu) g = g.to(device) use_cuda = True num_nodes = g.number_of_nodes() node_ids = tf.range(num_nodes, dtype=tf.int64) 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) target_idx = tf.squeeze(tf.where(loc)) # since the nodes are featureless, the input feature is then the node id. feats = tf.range(num_nodes, dtype=tf.int64) 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) # optimizer 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) 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)) # 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) t1 = time.time() grads = tape.gradient(loss, model.trainable_weights) optimizer.apply_gradients(zip(grads, model.trainable_weights)) t2 = time.time() 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])) train_acc = acc(logits, labels, train_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() 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_acc = acc(logits, labels, test_idx) 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:])))
def load_ogb(dataset, global_norm): if dataset == 'ogbn-mag': dataset = DglNodePropPredDataset(name=dataset) split_idx = dataset.get_idx_split() train_idx = split_idx["train"]['paper'] val_idx = split_idx["valid"]['paper'] test_idx = split_idx["test"]['paper'] hg_orig, labels = dataset[0] subgs = {} for etype in hg_orig.canonical_etypes: u, v = hg_orig.all_edges(etype=etype) subgs[etype] = (u, v) subgs[(etype[2], 'rev-' + etype[1], etype[0])] = (v, u) hg = dgl.heterograph(subgs) hg.nodes['paper'].data['feat'] = hg_orig.nodes['paper'].data['feat'] paper_labels = labels['paper'].squeeze() num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) num_classes = dataset.num_classes category = 'paper' print('Number of relations: {}'.format(num_rels)) print('Number of class: {}'.format(num_classes)) print('Number of train: {}'.format(len(train_idx))) print('Number of valid: {}'.format(len(val_idx))) print('Number of test: {}'.format(len(test_idx))) # currently we do not support node feature in mag dataset. # calculate norm for each edge type and store in edge if global_norm is False: for canonical_etype in hg.canonical_etypes: u, v, eid = hg.all_edges(form='all', etype=canonical_etype) _, 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_etype].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) if global_norm: u, v, eid = g.all_edges(form='all') _, 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) g.edata['norm'] = norm 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] train_idx = target_idx[train_idx] val_idx = target_idx[val_idx] test_idx = target_idx[test_idx] train_mask = th.zeros((g.number_of_nodes(), ), dtype=th.bool) train_mask[train_idx] = True val_mask = th.zeros((g.number_of_nodes(), ), dtype=th.bool) val_mask[val_idx] = True test_mask = th.zeros((g.number_of_nodes(), ), dtype=th.bool) test_mask[test_idx] = True g.ndata['train_mask'] = train_mask g.ndata['val_mask'] = val_mask g.ndata['test_mask'] = test_mask labels = th.full((g.number_of_nodes(), ), -1, dtype=paper_labels.dtype) labels[target_idx] = paper_labels g.ndata['labels'] = labels return g else: raise ("Do not support other ogbn datasets.")
def forward(self, g): x = g.nodes['atom'].data['h'] g_sub = dgl.to_homo(g.edge_type_subgraph(['atom_neighbors_atom'])) x = self.gn(g_sub, x) g.nodes['atom'].data['h'] = x return g
def get_unigraph (self): g = self.get_whole_graph() return graph.to_homo(g)