""" Modeling Relational Data with Graph Convolutional Networks Paper: https://arxiv.org/abs/1703.06103 Code: https://github.com/tkipf/relational-gcn Difference compared to tkipf/relation-gcn * l2norm applied to all weights * remove nodes that won't be touched """ import argparse, gc import numpy as np import time import torch as th import torch.nn as nn import dgl.function as fn import torch.nn.functional as F import dgl import dgl.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel from dgl import DGLGraph from functools import partial from dgl.data.rdf import AIFBDataset from src.skeleton.graph_builder import StandaloneGraphBuilder from src.skeleton.train_type import SamplingGraphTraining from src.application.rgcn.rgcn import RelGraphEmbedLayer, EntityClassify from dgl.contrib.hostmap_tensor import HostMapTensor from src.skeleton.dataloader import Dataloader import tqdm from sklearn.metrics import roc_auc_score # from torch.utils.tensorboard import SummaryWriter ''' 这是单机的异构图节点分类任务-Demo: 适用于: -- 图的数据量较大，比如100万~1亿点, 1000万~10亿边。 class RgcnGraphBuilder 负责加载数据 class RgcnTrainer 负责训练和预测 class RgcnTrainingDataLoader 负责做训练采样和数据遍历 用户如果需要改动只需要： 1、改动RgcnGraphBuilder.build_dataset 此方法负责从DGL图中分离训练数据、预测数据、测试数据 2、改动RgcnTrainer.train 此方法负责训练逻辑 3、改动RgcnTrainer.evaluate 此方法负责离线预测逻辑 4、改动RgcnTrainingDataLoader.init 此方法负责输出返回一个迭代遍历器、用于遍历数据集 这里使用AIFB数据集做精度对齐（epoch=50, batch_size=128） 社区aifb数据集节点分类测试集精度： Final Test Accuracy: 0.9250 | Test loss: 0.3929 平台aifb数据集节点分类测试集精度： Final Test Accuracy: 0.9250 | Test loss: 0.2953 ''' class RgcnGraphBuilder(StandaloneGraphBuilder): def build_dataset(self, g): hg = g # category = self.flags.category num_classes = self.flags.num_classes num_rels = len(hg.canonical_etypes) num_of_ntype = len(hg.ntypes) # train_mask = hg.nodes[category].data.pop(self.flags.train_mask) # test_mask = hg.nodes[category].data.pop(self.flags.test_mask) # labels = hg.nodes[category].data.pop(self.flags.label) eids = th.arange(g.number_of_edges()) #eids = np.random.permutation(eids) val_size = int(len(eids) * 0.1) test_size = int(len(eids) * 0.2) # train_size = g.number_of_edges() - val_size - test_size # valid_eids = eids[:val_size] # test_eids = eids[val_size: val_size + test_size] # train_eids = eids[val_size + test_size:] valid_eids = dgl.contrib.HostMapTensor('valid_eids', eids[:val_size]) test_eids = dgl.contrib.HostMapTensor('test_eids', eids[val_size: val_size + test_size]) train_eids = dgl.contrib.HostMapTensor('train_eids', eids[val_size + test_size:]) # train_idx = th.nonzero(train_mask, as_tuple=False).squeeze() # test_idx = th.nonzero(test_mask, as_tuple=False).squeeze() # val_idx = train_idx node_feats = {} for ntype in hg.ntypes: if len(hg.nodes[ntype].data) == 0 or self.flags.node_feats is False: node_feats[str(hg.get_ntype_id(ntype))] = hg.number_of_nodes(ntype) else: assert len(hg.nodes[ntype].data) == 1 feat = hg.nodes[ntype].data.pop(self.flags.feat) if feat is not None: feats = HostMapTensor(ntype + '__' + self.flags.feat, feat) node_feats[str(hg.get_ntype_id(ntype))] = feats # get target category id # category_id = len(hg.ntypes) # for i, ntype in enumerate(hg.ntypes): # if ntype == category: # category_id = i # print('{}:{}'.format(i, ntype)) g = dgl.to_homogeneous(hg) ntype_tensor = g.ndata[dgl.NTYPE] ntype_tensor.share_memory_() etype_tensor = g.edata[dgl.ETYPE] etype_tensor = dgl.contrib.HostMapTensor('etype_tensor', etype_tensor) typeid_tensor = g.ndata[dgl.NID] typeid_tensor.share_memory_() #ntype_tensor = dgl.contrib.HostMapTensor('ntype_tensor', g.ndata[dgl.NTYPE]) #etype_tensor = dgl.contrib.HostMapTensor('etype_tensor', g.edata[dgl.ETYPE]) #typeid_tensor = dgl.contrib.HostMapTensor('typeid_tensor', g.edata[dgl.NID]) # 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_() # # This is a graph with multiple node types, so we want a way to map # # our target node from their global node numberings, back to their # # numberings within their type. This is used when taking the nodes in a # # mini-batch, and looking up their type-specific labels # inv_target = th.empty(node_ids.shape, # dtype=node_ids.dtype) # inv_target.share_memory_() # inv_target[target_idx] = th.arange(0, target_idx.shape[0], # dtype=inv_target.dtype) # Create csr/coo/csc formats before launching training processes with multi-gpu. # This avoids creating certain formats in each sub-process, which saves momory and CPU. g.create_formats_() g = g.shared_memory('g') return g, node_feats, num_of_ntype, num_classes, num_rels, train_eids, valid_eids, test_eids, ntype_tensor, etype_tensor, typeid_tensor class RgcnTrainer(SamplingGraphTraining): def train(self, g, dataset, device, n_gpus, proc_id, **kwargs): dev_id = -1 if n_gpus == 0 else device.index queue = kwargs['queue'] if n_gpus > 1 else None g, node_feats, num_of_ntype, num_classes, num_rels, train_eids, valid_eids, test_eids, ntype_tensor, etype_tensor, typeid_tensor = dataset node_tids = ntype_tensor world_size = n_gpus if n_gpus > 0: etype_tensor.uva(device) for key in node_feats: if not isinstance(node_feats[key], int): node_feats[key].uva(device) if n_gpus == 1: g = g.to(device) if n_gpus > 1: g = g.uva(device) dist_init_method = 'tcp://{master_ip}:{master_port}'.format( master_ip='127.0.0.1', master_port=self.flags.master_port) th.distributed.init_process_group(backend=self.flags.communication_backend, init_method=dist_init_method, world_size=world_size, rank=proc_id) # node features # None for one-hot feature, if not none, it should be the feature tensor. embed_layer = RelGraphEmbedLayer(dev_id if self.flags.embedding_gpu or not self.flags.dgl_sparse else -1, dev_id, g.number_of_nodes(), node_tids, num_of_ntype, node_feats, self.flags.num_hidden, dgl_sparse=self.flags.dgl_sparse) # 设置目标函数 loss_fcn = CrossEntropyLoss() # create model # all model params are in device. model = EntityClassify(dev_id, g.number_of_nodes(), self.flags.num_hidden, num_classes, num_rels, num_bases=self.flags.num_bases, num_hidden_layers=self.flags.num_layers - 2, dropout=self.flags.dropout, use_self_loop=self.flags.use_self_loop, low_mem=self.flags.low_mem, layer_norm=self.flags.layer_norm) if n_gpus == 1: th.cuda.set_device(dev_id) model.cuda(dev_id) if self.flags.dgl_sparse: embed_layer.cuda(dev_id) elif n_gpus > 1: if dev_id >= 0: model.cuda(dev_id) model = DistributedDataParallel(model, device_ids=[dev_id], output_device=dev_id) if self.flags.dgl_sparse: embed_layer.cuda(dev_id) if len(list(embed_layer.parameters())) > 0: embed_layer = DistributedDataParallel(embed_layer, device_ids=[dev_id], output_device=dev_id) else: if len(list(embed_layer.parameters())) > 0: embed_layer = DistributedDataParallel(embed_layer, device_ids=None, output_device=None) # optimizer dense_params = list(model.parameters()) if self.flags.node_feats: if n_gpus > 1: dense_params += list(embed_layer.module.embeds.parameters()) else: dense_params += list(embed_layer.embeds.parameters()) optimizer = th.optim.Adam(dense_params, lr=self.flags.lr, weight_decay=self.flags.l2norm) if self.flags.dgl_sparse: all_params = list(model.parameters()) + list(embed_layer.parameters()) optimizer = th.optim.Adam(all_params, lr=self.flags.lr, weight_decay=self.flags.l2norm) if n_gpus > 1 and isinstance(embed_layer, DistributedDataParallel): dgl_emb = embed_layer.module.dgl_emb else: dgl_emb = embed_layer.dgl_emb emb_optimizer = dgl.optim.SparseAdam(params=dgl_emb, lr=self.flags.sparse_lr, eps=1e-8) if len(dgl_emb) > 0 else None else: if n_gpus > 1: embs = list(embed_layer.module.node_embeds.parameters()) else: embs = list(embed_layer.node_embeds.parameters()) emb_optimizer = th.optim.SparseAdam(embs, lr=self.flags.sparse_lr) if len(embs) > 0 else None ntype_tensor = ntype_tensor.to(device) # etype_tensor = etype_tensor.to(device) typeid_tensor = typeid_tensor.to(device) # train_eids = train_eids.to(device) # valid_eids = valid_eids.to(device) # test_eids = test_eids.to(device) dataset = train_eids, valid_eids, test_eids, device dataloader = RgcnTrainingDataLoader(self.flags).init(g, dataset) loader, val_loader, test_loader = dataloader # training loop print("start training...") forward_time = [] backward_time = [] train_time = 0 validation_time = 0 test_time = 0 last_val_acc = 0.0 do_test = False for epoch in range(self.flags.num_epochs): if n_gpus > 1: loader.set_epoch(epoch) tstart = time.time() model.train() embed_layer.train() # for i, sample_data in enumerate(loader): for i, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(loader): # input_nodes, seeds, blocks = sample_data # # map the seed nodes back to their type-specific ids, so that they # # can be used to look up their respective labels # seeds = inv_target[seeds] for block in blocks: gen_norm(block, ntype_tensor, etype_tensor, typeid_tensor) t0 = time.time() feats = embed_layer(blocks[0].srcdata[dgl.NID], blocks[0].srcdata['ntype'], blocks[0].srcdata['type_id'], node_feats) blocks = [block.long().to(device) for block in blocks] # logits = model(blocks, feats) pos_graph = pos_graph.to(device) neg_graph = neg_graph.to(device) batch_pred = model(blocks, feats) f_step = time.time() loss = loss_fcn(batch_pred, pos_graph, neg_graph) # loss = F.cross_entropy(logits, labels[seeds]) # writer.add_scalar('loss', loss, global_step) t1 = time.time() optimizer.zero_grad() if emb_optimizer is not None: emb_optimizer.zero_grad() loss.backward() if emb_optimizer is not None: emb_optimizer.step() optimizer.step() t2 = time.time() forward_time.append(t1 - t0) backward_time.append(t2 - t1) # train_acc = th.sum(logits.argmax(dim=1) == labels[seeds]).item() / len(seeds) if i % 100 == 0 and proc_id == 0: print("Train Loss: {:.4f}". format(loss.item())) # writer.add_scalar('train_acc', train_acc, global_step) # global_step += 1 print("Epoch {:05d}:{:05d} | Train Forward Time(s) {:.4f} | Backward Time(s) {:.4f}". format(epoch, self.flags.num_epochs, forward_time[-1], backward_time[-1])) tend = time.time() train_time += (tend - tstart) # val_acc, val_loss, validation_time = self._evaluate(n_gpus, labels, queue, proc_id, model, embed_layer, # val_loader, node_feats, inv_target, 'Validation') # do_test = val_acc > last_val_acc # last_val_acc = val_acc # if n_gpus > 1: # th.distributed.barrier() # if proc_id == 0: # for i in range(1, n_gpus): # queue.put(do_test) # else: # do_test = queue.get() # if epoch == self.flags.num_epochs - 1 or (epoch > 0 and do_test): # test_acc, test_loss, test_time = self._evaluate(n_gpus, labels, queue, proc_id, model, embed_layer, # test_loader, node_feats, inv_target, 'Test') # if n_gpus > 1: # th.distributed.barrier() print("{}/{} Mean forward time: {:4f}".format(proc_id, n_gpus, np.mean(forward_time[len(forward_time) // 4:]))) print("{}/{} Mean backward time: {:4f}".format(proc_id, n_gpus, np.mean(backward_time[len(backward_time) // 4:]))) # if proc_id == 0: # print("Final Test Accuracy: {:.4f} | Test loss: {:.4f}".format(test_acc, test_loss)) # print("Train {}s, valid {}s, test {}s".format(train_time, validation_time, test_time)) def _evaluate(self, n_gpus, labels, queue, proc_id, model, embed_layer, data_loader, node_feats, inv_target, mode): tstart = time.time() time_cost = 0 acc = 0 loss = 0 logits, seeds = evaluate(model, embed_layer, data_loader, node_feats, inv_target) if queue is not None: queue.put((logits, seeds)) if proc_id == 0: loss, acc = self._collect_eval(n_gpus, labels, queue) if queue is not None else \ (F.cross_entropy(logits, labels[seeds].cpu()).item(), \ th.sum(logits.argmax(dim=1) == labels[seeds].cpu()).item() / len(seeds)) print("{} Accuracy: {:.4f} | {} loss: {:.4f}".format(mode, acc, mode, loss)) tend = time.time() time_cost = (tend-tstart) return acc, loss, time_cost def _collect_eval(self, n_gpus, labels, queue): eval_logits = [] eval_seeds = [] for i in range(n_gpus): log = queue.get() eval_l, eval_s = log eval_logits.append(eval_l) eval_seeds.append(eval_s) eval_logits = th.cat(eval_logits) eval_seeds = th.cat(eval_seeds) eval_loss = F.cross_entropy(eval_logits, labels[eval_seeds].cpu()).item() eval_acc = th.sum(eval_logits.argmax(dim=1) == labels[eval_seeds].cpu()).item() / len(eval_seeds) return eval_loss, eval_acc class RgcnTrainingDataLoader(Dataloader): def init(self, g, dataset): train_eids, valid_eids, test_eids, device = dataset # target_idx = target_idx.to(device) # 查找有几块GPU n_gpus = len(list(map(int, self.flags.gpu.split(',')))) # 每层邻居数 fanouts = [int(fanout) for fanout in self.flags.fanout.split(',')] sampler = dgl.dataloading.MultiLayerNeighborSampler(fanouts) loader = dgl.dataloading.EdgeDataLoader( g, train_eids, sampler, negative_sampler=dgl.dataloading.negative_sampler.Uniform(5), batch_size=self.flags.batch_size, device=device, use_ddp=n_gpus > 1, shuffle=True, drop_last=False, num_workers=self.flags.num_workers) val_loader = dgl.dataloading.EdgeDataLoader( g, valid_eids, sampler, negative_sampler=dgl.dataloading.negative_sampler.Uniform(5), batch_size=self.flags.batch_size, device=device, use_ddp=n_gpus > 1, shuffle=False, drop_last=False, num_workers=self.flags.num_workers) test_loader = dgl.dataloading.EdgeDataLoader( g, test_eids, sampler, negative_sampler=dgl.dataloading.negative_sampler.Uniform(5), batch_size=self.flags.batch_size, device=device, use_ddp=n_gpus > 1, shuffle=True, drop_last=False, num_workers=self.flags.num_workers) # loader = dgl.dataloading.NodeDataLoader( # g, # target_idx[train_idx], # sampler, # use_ddp=n_gpus > 1, # device=device if self.flags.num_workers == 0 else None, # batch_size=self.flags.batch_size, # shuffle=True, # drop_last=False, # num_workers=self.flags.num_workers) # # validation sampler # val_loader = dgl.dataloading.NodeDataLoader( # g, # target_idx[val_idx], # sampler, # use_ddp=n_gpus > 1, # device=device if self.flags.num_workers == 0 else None, # batch_size=self.flags.batch_size, # shuffle=False, # drop_last=False, # num_workers=self.flags.num_workers) # # test sampler # test_sampler = dgl.dataloading.MultiLayerNeighborSampler([-1] * self.flags.num_layers) # test_loader = dgl.dataloading.NodeDataLoader( # g, # target_idx[test_idx], # test_sampler, # use_ddp=n_gpus > 1, # device=device if self.flags.num_workers == 0 else None, # batch_size=self.flags.eval_batch_size, # shuffle=False, # drop_last=False, # num_workers=self.flags.num_workers) return loader, val_loader, test_loader def gen_norm(g, ntype_tensor, etype_tensor, typeid_tensor): _, 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], device=eid.device) / degrees norm = norm.unsqueeze(1) g.edata['norm'] = norm g.srcdata['ntype'] = ntype_tensor[g.srcdata[dgl.NID]] g.edata['etype'] = etype_tensor[eid] g.srcdata['type_id'] = typeid_tensor[g.srcdata[dgl.NID]] def evaluate(model, embed_layer, eval_loader, node_feats, inv_target, ntype_tensor, etype_tensor, typeid_tensor): model.eval() embed_layer.eval() eval_logits = [] eval_seeds = [] with th.no_grad(): th.cuda.empty_cache() for i, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(eval_loader): for block in blocks: gen_norm(block, ntype_tensor, etype_tensor, typeid_tensor) feats = embed_layer(blocks[0].srcdata[dgl.NID], blocks[0].srcdata['ntype'], blocks[0].srcdata['type_id'], node_feats) logits = model(blocks, feats) loss_fcn = AUC() auc = loss_fcn(logits, pos_graph, neg_graph) print("valid auc: {:.4f}". format(auc.item())) # eval_logits.append(logits.cpu()) eval_logits = th.cat(eval_logits) eval_seeds = th.cat(eval_seeds) return eval_logits, eval_seeds class CrossEntropyLoss(nn.Module): def forward(self, block_outputs, pos_graph, neg_graph): with pos_graph.local_scope(): pos_graph.ndata['h'] = block_outputs pos_graph.apply_edges(fn.u_dot_v('h', 'h', 'score')) pos_score = pos_graph.edata['score'] with neg_graph.local_scope(): neg_graph.ndata['h'] = block_outputs neg_graph.apply_edges(fn.u_dot_v('h', 'h', 'score')) neg_score = neg_graph.edata['score'] score = th.cat([pos_score, neg_score]) label = th.cat([th.ones_like(pos_score), th.zeros_like(neg_score)]).long() loss = F.binary_cross_entropy_with_logits(score, label.float()) return loss class AUC(nn.Module): def forward(self, block_outputs, pos_graph, neg_graph): with pos_graph.local_scope(): pos_graph.ndata['h'] = block_outputs pos_graph.apply_edges(fn.u_dot_v('h', 'h', 'score')) pos_score = pos_graph.edata['score'] with neg_graph.local_scope(): neg_graph.ndata['h'] = block_outputs neg_graph.apply_edges(fn.u_dot_v('h', 'h', 'score')) neg_score = neg_graph.edata['score'] score = th.cat([pos_score, neg_score]).numpy() label = th.cat([th.ones_like(pos_score), th.zeros_like(neg_score)]).numpy() return roc_auc_score(label, score)