def __init__(self, hidden_size): super(LoopyBPUpdate, self).__init__() self.hidden_size = hidden_size self.W_h = nn.Linear(hidden_size, hidden_size, bias=False) def forward(self, node): msg_input = node.data['msg_input'] msg_delta = self.W_h(node.data['accum_msg'] + node.data['alpha']) msg = torch.relu(msg_input + msg_delta) return {'msg': msg} if PAPER: mpn_gather_msg = [ DGLF.copy_edge(edge='msg', out='msg'), DGLF.copy_edge(edge='alpha', out='alpha') ] else: mpn_gather_msg = DGLF.copy_edge(edge='msg', out='msg') if PAPER: mpn_gather_reduce = [ DGLF.sum(msg='msg', out='m'), DGLF.sum(msg='alpha', out='accum_alpha'), ] else: mpn_gather_reduce = DGLF.sum(msg='msg', out='m') class GatherUpdate(nn.Module):
edges = bfs_edges_generator(forest, roots) _, leaves = forest.find_edges(edges[-1]) edges_back = bfs_edges_generator(forest, roots, reverse=True) yield from reversed(edges_back) yield from edges enc_tree_msg = [ DGLF.copy_src(src='m', out='m'), DGLF.copy_src(src='rm', out='rm') ] enc_tree_reduce = [ DGLF.sum(msg='m', out='s'), DGLF.sum(msg='rm', out='accum_rm') ] enc_tree_gather_msg = DGLF.copy_edge(edge='m', out='m') enc_tree_gather_reduce = DGLF.sum(msg='m', out='m') class EncoderGatherUpdate(nn.Module): def __init__(self, hidden_size): nn.Module.__init__(self) self.hidden_size = hidden_size self.W = nn.Linear(2 * hidden_size, hidden_size) def reset_parameters(self): """Reinitialize model parameters.""" self.W.reset_parameters() def forward(self, nodes):
def check_partition(g, part_method, reshuffle): g.ndata['labels'] = F.arange(0, g.number_of_nodes()) g.ndata['feats'] = F.tensor(np.random.randn(g.number_of_nodes(), 10), F.float32) g.edata['feats'] = F.tensor(np.random.randn(g.number_of_edges(), 10), F.float32) g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h')) g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh')) num_parts = 4 num_hops = 2 orig_nids, orig_eids = partition_graph(g, 'test', num_parts, '/tmp/partition', num_hops=num_hops, part_method=part_method, reshuffle=reshuffle, return_mapping=True) part_sizes = [] for i in range(num_parts): part_g, node_feats, edge_feats, gpb, _, ntypes, etypes = load_partition( '/tmp/partition/test.json', i) # Check the metadata assert gpb._num_nodes() == g.number_of_nodes() assert gpb._num_edges() == g.number_of_edges() assert gpb.num_partitions() == num_parts gpb_meta = gpb.metadata() assert len(gpb_meta) == num_parts assert len(gpb.partid2nids(i)) == gpb_meta[i]['num_nodes'] assert len(gpb.partid2eids(i)) == gpb_meta[i]['num_edges'] part_sizes.append((gpb_meta[i]['num_nodes'], gpb_meta[i]['num_edges'])) local_nid = gpb.nid2localnid( F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']), i) assert F.dtype(local_nid) in (F.int64, F.int32) assert np.all(F.asnumpy(local_nid) == np.arange(0, len(local_nid))) local_eid = gpb.eid2localeid( F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']), i) assert F.dtype(local_eid) in (F.int64, F.int32) assert np.all(F.asnumpy(local_eid) == np.arange(0, len(local_eid))) # Check the node map. local_nodes = F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']) llocal_nodes = F.nonzero_1d(part_g.ndata['inner_node']) local_nodes1 = gpb.partid2nids(i) assert F.dtype(local_nodes1) in (F.int32, F.int64) assert np.all( np.sort(F.asnumpy(local_nodes)) == np.sort(F.asnumpy( local_nodes1))) # Check the edge map. local_edges = F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']) local_edges1 = gpb.partid2eids(i) assert F.dtype(local_edges1) in (F.int32, F.int64) assert np.all( np.sort(F.asnumpy(local_edges)) == np.sort(F.asnumpy( local_edges1))) # Verify the mapping between the reshuffled IDs and the original IDs. part_src_ids, part_dst_ids = part_g.edges() part_src_ids = F.gather_row(part_g.ndata[dgl.NID], part_src_ids) part_dst_ids = F.gather_row(part_g.ndata[dgl.NID], part_dst_ids) part_eids = part_g.edata[dgl.EID] orig_src_ids = F.gather_row(orig_nids, part_src_ids) orig_dst_ids = F.gather_row(orig_nids, part_dst_ids) orig_eids1 = F.gather_row(orig_eids, part_eids) orig_eids2 = g.edge_ids(orig_src_ids, orig_dst_ids) assert F.shape(orig_eids1)[0] == F.shape(orig_eids2)[0] assert np.all(F.asnumpy(orig_eids1) == F.asnumpy(orig_eids2)) if reshuffle: part_g.ndata['feats'] = F.gather_row(g.ndata['feats'], part_g.ndata['orig_id']) part_g.edata['feats'] = F.gather_row(g.edata['feats'], part_g.edata['orig_id']) # when we read node data from the original global graph, we should use orig_id. local_nodes = F.boolean_mask(part_g.ndata['orig_id'], part_g.ndata['inner_node']) local_edges = F.boolean_mask(part_g.edata['orig_id'], part_g.edata['inner_edge']) else: part_g.ndata['feats'] = F.gather_row(g.ndata['feats'], part_g.ndata[dgl.NID]) part_g.edata['feats'] = F.gather_row(g.edata['feats'], part_g.edata[dgl.NID]) part_g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h')) part_g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh')) assert F.allclose(F.gather_row(g.ndata['h'], local_nodes), F.gather_row(part_g.ndata['h'], llocal_nodes)) assert F.allclose(F.gather_row(g.ndata['eh'], local_nodes), F.gather_row(part_g.ndata['eh'], llocal_nodes)) for name in ['labels', 'feats']: assert '_N/' + name in node_feats assert node_feats['_N/' + name].shape[0] == len(local_nodes) assert np.all( F.asnumpy(g.ndata[name])[F.asnumpy(local_nodes)] == F.asnumpy( node_feats['_N/' + name])) for name in ['feats']: assert '_E/' + name in edge_feats assert edge_feats['_E/' + name].shape[0] == len(local_edges) assert np.all( F.asnumpy(g.edata[name])[F.asnumpy(local_edges)] == F.asnumpy( edge_feats['_E/' + name])) if reshuffle: node_map = [] edge_map = [] for i, (num_nodes, num_edges) in enumerate(part_sizes): node_map.append(np.ones(num_nodes) * i) edge_map.append(np.ones(num_edges) * i) node_map = np.concatenate(node_map) edge_map = np.concatenate(edge_map) nid2pid = gpb.nid2partid(F.arange(0, len(node_map))) assert F.dtype(nid2pid) in (F.int32, F.int64) assert np.all(F.asnumpy(nid2pid) == node_map) eid2pid = gpb.eid2partid(F.arange(0, len(edge_map))) assert F.dtype(eid2pid) in (F.int32, F.int64) assert np.all(F.asnumpy(eid2pid) == edge_map)
MAX_NB = 8 MAX_DECODE_LEN = 100 def dfs_order(forest, roots): forest = tocpu(forest) edges = dfs_labeled_edges_generator(forest, roots, has_reverse_edge=True) for e, l in zip(*edges): # I exploited the fact that the reverse edge ID equal to 1 xor forward # edge ID for molecule trees. Normally, I should locate reverse edges # using find_edges(). yield e ^ l, l dec_tree_node_msg = DGLF.copy_edge(edge='m', out='m') dec_tree_node_reduce = DGLF.sum(msg='m', out='h') def dec_tree_node_update(nodes): return {'new': nodes.data['new'].clone().zero_()} dec_tree_edge_msg = [ DGLF.copy_src(src='m', out='m'), DGLF.copy_src(src='rm', out='rm') ] dec_tree_edge_reduce = [ DGLF.sum(msg='m', out='s'), DGLF.sum(msg='rm', out='accum_rm') ]
def forward(self, g, node_feats, edge_feats, node_only=False): r"""Update node and edge representations. Parameters ---------- g : DGLGraph DGLGraph for a batch of graphs node_feats : float32 tensor of shape (V, node_in_feats) Input node features. V for the number of nodes in the batch of graphs. edge_feats : float32 tensor of shape (E, edge_in_feats) Input edge features. E for the number of edges in the batch of graphs. node_only : bool Whether to update node representations only. If False, edge representations will be updated as well. Default to False. Returns ------- new_node_feats : float32 tensor of shape (V, node_out_feats) Updated node representations. new_edge_feats : float32 tensor of shape (E, edge_out_feats) Updated edge representations. """ g = g.local_var() ####### g.srcdata['h'] = node_feats g.ndata['feat'] = node_feats g.apply_edges( lambda edges: { 'e': torch.sum( (torch.mul(edges.src['h'], torch.tanh(edges.dst['h']))), 1) }) e = self.leaky_relu(g.edata.pop('e')) e_soft = edge_softmax(g, e) g.ndata.pop('feat') g.srcdata.pop('h') #print(e_soft.shape) ####### # Update node features node_node_feats = self.activation( self.node_to_node(node_feats)) # torch.Size([596, 50]) g.edata['e2n'] = self.activation(self.edge_to_node(edge_feats)) g.update_all(fn.copy_edge('e2n', 'm'), fn.sum('m', 'e2n')) edge_node_feats = g.ndata.pop('e2n') # torch.Size([596, 50]) new_node_feats = self.activation( self.update_node( torch.cat([node_node_feats, edge_node_feats], dim=1))) # torch.Size([596, 50]) if node_only: return new_node_feats # Update edge features g.ndata['left_hv'] = self.left_node_to_edge(node_feats) g.ndata['right_hv'] = self.right_node_to_edge(node_feats) g.apply_edges(fn.u_add_v('left_hv', 'right_hv', 'first')) g.apply_edges(fn.u_add_v('right_hv', 'left_hv', 'second')) first_edge_feats = self.activation(g.edata.pop('first')) second_edge_feats = self.activation(g.edata.pop('second')) third_edge_feats = self.activation(self.edge_to_edge(edge_feats)) new_edge_feats = self.activation( self.update_edge( torch.cat( [first_edge_feats, second_edge_feats, third_edge_feats], dim=1))) return new_node_feats, new_edge_feats, e_soft
class LoopyBPUpdate(nn.Module): def __init__(self, hidden_size): super(LoopyBPUpdate, self).__init__() self.hidden_size = hidden_size self.W_h = nn.Linear(hidden_size, hidden_size, bias=False) def forward(self, nodes): msg_input = nodes.data['msg_input'] msg_delta = self.W_h(nodes.data['accum_msg']) msg = F.relu(msg_input + msg_delta) return {'msg': msg} mpn_gather_msg = DGLF.copy_edge(edge='msg', out='msg') mpn_gather_reduce = DGLF.sum(msg='msg', out='m') class GatherUpdate(nn.Module): def __init__(self, hidden_size): super(GatherUpdate, self).__init__() self.hidden_size = hidden_size self.W_o = nn.Linear(ATOM_FDIM + hidden_size, hidden_size) def forward(self, nodes): m = nodes.data['m'] return { 'h': F.relu(self.W_o(torch.cat([nodes.data['x'], m], 1))), }
self.W_h = nn.Linear(hidden_size, hidden_size, bias=False) def reset_parameters(self): """Reinitialize model parameters.""" self.W_h.reset_parameters() def forward(self, node): msg_input = node.data['msg_input'] msg_delta = self.W_h(node.data['accum_msg'] + node.data['alpha']) msg = torch.relu(msg_input + msg_delta) return {'msg': msg} if PAPER: mpn_gather_msg = [ fn.copy_edge(edge='msg', out='msg'), fn.copy_edge(edge='alpha', out='alpha') ] else: mpn_gather_msg = fn.copy_edge(edge='msg', out='msg') if PAPER: mpn_gather_reduce = [ fn.sum(msg='msg', out='m'), fn.sum(msg='alpha', out='accum_alpha'), ] else: mpn_gather_reduce = fn.sum(msg='msg', out='m') class GatherUpdate(nn.Module): def __init__(self, hidden_size):
def check_partition(reshuffle): g = create_random_graph(10000) g.ndata['labels'] = F.arange(0, g.number_of_nodes()) g.ndata['feats'] = F.tensor(np.random.randn(g.number_of_nodes(), 10)) g.edata['feats'] = F.tensor(np.random.randn(g.number_of_edges(), 10)) g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h')) g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh')) num_parts = 4 num_hops = 2 partition_graph(g, 'test', num_parts, '/tmp/partition', num_hops=num_hops, part_method='metis', reshuffle=reshuffle) part_sizes = [] for i in range(num_parts): part_g, node_feats, edge_feats, gpb = load_partition( '/tmp/partition/test.json', i) # Check the metadata assert gpb._num_nodes() == g.number_of_nodes() assert gpb._num_edges() == g.number_of_edges() assert gpb.num_partitions() == num_parts gpb_meta = gpb.metadata() assert len(gpb_meta) == num_parts assert len(gpb.partid2nids(i)) == gpb_meta[i]['num_nodes'] assert len(gpb.partid2eids(i)) == gpb_meta[i]['num_edges'] part_sizes.append((gpb_meta[i]['num_nodes'], gpb_meta[i]['num_edges'])) local_nid = gpb.nid2localnid( F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']), i) assert np.all(F.asnumpy(local_nid) == np.arange(0, len(local_nid))) local_eid = gpb.eid2localeid( F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']), i) assert np.all(F.asnumpy(local_eid) == np.arange(0, len(local_eid))) # Check the node map. local_nodes = F.boolean_mask(part_g.ndata[dgl.NID], part_g.ndata['inner_node']) llocal_nodes = F.nonzero_1d(part_g.ndata['inner_node']) local_nodes1 = gpb.partid2nids(i) assert np.all( np.sort(F.asnumpy(local_nodes)) == np.sort(F.asnumpy( local_nodes1))) # Check the edge map. local_edges = F.boolean_mask(part_g.edata[dgl.EID], part_g.edata['inner_edge']) local_edges1 = gpb.partid2eids(i) assert np.all( np.sort(F.asnumpy(local_edges)) == np.sort(F.asnumpy( local_edges1))) if reshuffle: part_g.ndata['feats'] = F.gather_row(g.ndata['feats'], part_g.ndata['orig_id']) part_g.edata['feats'] = F.gather_row(g.edata['feats'], part_g.edata['orig_id']) # when we read node data from the original global graph, we should use orig_id. local_nodes = F.boolean_mask(part_g.ndata['orig_id'], part_g.ndata['inner_node']) local_edges = F.boolean_mask(part_g.edata['orig_id'], part_g.edata['inner_edge']) else: part_g.ndata['feats'] = F.gather_row(g.ndata['feats'], part_g.ndata[dgl.NID]) part_g.edata['feats'] = F.gather_row(g.edata['feats'], part_g.edata[dgl.NID]) part_g.update_all(fn.copy_src('feats', 'msg'), fn.sum('msg', 'h')) part_g.update_all(fn.copy_edge('feats', 'msg'), fn.sum('msg', 'eh')) assert F.allclose(F.gather_row(g.ndata['h'], local_nodes), F.gather_row(part_g.ndata['h'], llocal_nodes)) assert F.allclose(F.gather_row(g.ndata['eh'], local_nodes), F.gather_row(part_g.ndata['eh'], llocal_nodes)) for name in ['labels', 'feats']: assert name in node_feats assert node_feats[name].shape[0] == len(local_nodes) assert np.all( F.asnumpy(g.ndata[name])[F.asnumpy(local_nodes)] == F.asnumpy( node_feats[name])) for name in ['feats']: assert name in edge_feats assert edge_feats[name].shape[0] == len(local_edges) assert np.all( F.asnumpy(g.edata[name])[F.asnumpy(local_edges)] == F.asnumpy( edge_feats[name])) if reshuffle: node_map = [] edge_map = [] for i, (num_nodes, num_edges) in enumerate(part_sizes): node_map.append(np.ones(num_nodes) * i) edge_map.append(np.ones(num_edges) * i) node_map = np.concatenate(node_map) edge_map = np.concatenate(edge_map) assert np.all( F.asnumpy(gpb.nid2partid(F.arange(0, len(node_map)))) == node_map) assert np.all( F.asnumpy(gpb.eid2partid(F.arange(0, len(edge_map)))) == edge_map)
def forward(self, graph, feat, lambda_max=None): r""" Description ----------- Compute ChebNet layer. Parameters ---------- graph : DGLGraph The graph. feat : torch.Tensor The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes. lambda_max : list or tensor or None, optional. A list(tensor) with length :math:`B`, stores the largest eigenvalue of the normalized laplacian of each individual graph in ``graph``, where :math:`B` is the batch size of the input graph. Default: None. If None, this method would compute the list by calling ``dgl.laplacian_lambda_max``. Returns ------- torch.Tensor The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}` is size of output feature. """ def unnLaplacian(feat, D_invsqrt, graph): """ Operation Feat * D^-1/2 A D^-1/2 但是如果写成矩阵乘法:D^-1/2 A D^-1/2 Feat""" graph.ndata['h'] = feat * D_invsqrt graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h')) return graph.ndata.pop('h') * D_invsqrt with graph.local_scope(): #一点修改,这是原来的代码 if self.is_mnist: graph.update_all(fn.copy_edge('v', 'm'), fn.sum('m', 'h')) # 'v'与coordinate.py有关 D_invsqrt = th.pow( graph.ndata.pop('h').float().clamp(min=1), -0.5).unsqueeze(-1).to(feat.device) #D_invsqrt = th.pow(graph.in_degrees().float().clamp( # min=1), -0.5).unsqueeze(-1).to(feat.device) #print("in_degree : ",graph.in_degrees().shape) else: D_invsqrt = th.pow(graph.in_degrees().float().clamp(min=1), -0.5).unsqueeze(-1).to(feat.device) #print("D_invsqrt : ",D_invsqrt.shape) #print("ndata : ",graph.ndata['h'].shape) if lambda_max is None: try: lambda_max = laplacian_lambda_max(graph) except BaseException: # if the largest eigenvalue is not found dgl_warning( "Largest eigonvalue not found, using default value 2 for lambda_max", RuntimeWarning) lambda_max = th.Tensor(2).to(feat.device) if isinstance(lambda_max, list): lambda_max = th.Tensor(lambda_max).to(feat.device) if lambda_max.dim() == 1: lambda_max = lambda_max.unsqueeze(-1) # (B,) to (B, 1) # broadcast from (B, 1) to (N, 1) lambda_max = broadcast_nodes(graph, lambda_max) re_norm = 2. / lambda_max # X_0 is the raw feature, Xt refers to the concatenation of X_0, X_1, ... X_t Xt = X_0 = feat # X_1(f) if self._k > 1: h = unnLaplacian(X_0, D_invsqrt, graph) X_1 = -re_norm * h + X_0 * (re_norm - 1) # Concatenate Xt and X_1 Xt = th.cat((Xt, X_1), 1) # Xi(x), i = 2...k for _ in range(2, self._k): h = unnLaplacian(X_1, D_invsqrt, graph) X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0 # Concatenate Xt and X_i Xt = th.cat((Xt, X_i), 1) X_1, X_0 = X_i, X_1 # linear projection h = self.linear(Xt) # activation if self.activation: h = self.activation(h) #print('ChebConv.py Line163 h : ',h.shape) return h