|
- # The implementation of ICDM 2019 paper "Relation Structure-Aware Heterogeneous Graph Neural Network" RSHN.
- # @Time : 2021/3/1
- # @Author : Tianyu Zhao
- # @Email : tyzhao@bupt.edu.cn
- import dgl
- import torch as th
- import torch.nn as nn
- import torch.nn.functional as F
- from dgl import function as fn
- from dgl.utils import expand_as_pair
- from dgl.nn.functional import edge_softmax
- from . import BaseModel, register_model
- from ..sampler.RSHN_sampler import coarsened_line_graph
-
-
- @register_model('RSHN')
- class RSHN(BaseModel):
- r"""
- Relation structure-aware heterogeneous graph neural network (RSHN) builds coarsened line graph to obtain edge features first,
- then uses a novel Message Passing Neural Network (MPNN) to propagate node and edge features.
-
- We implement a API build a coarsened line graph.
-
- Attributes
- -----------
- edge_layers : AGNNConv
- Applied in Edge Layer.
- coarsened line graph : dgl.DGLGraph
- Propagate edge features.
- """
-
- @classmethod
- def build_model_from_args(cls, args, hg):
- rshn = cls(dim=args.hidden_dim,
- out_dim=args.out_dim,
- num_node_layer=args.num_node_layer,
- num_edge_layer=args.num_edge_layer,
- dropout=args.dropout
- )
- cl = coarsened_line_graph(rw_len=args.rw_len, batch_size=args.batch_size, n_dataset=args.dataset,
- symmetric=True)
- cl_graph = cl.get_cl_graph(hg).to(args.device)
- cl_graph = cl.init_cl_graph(cl_graph)
- rshn.cl_graph = cl_graph
-
- linear_e1 = nn.Linear(in_features=cl_graph.num_nodes(), out_features=args.hidden_dim, bias=False)
- nn.init.xavier_uniform_(linear_e1.weight)
- rshn.linear_e1 = linear_e1
- return rshn
-
- def __init__(self, dim, out_dim, num_node_layer, num_edge_layer, dropout):
- super(RSHN, self).__init__()
- # map the edge feature
- self.num_node_layer = num_node_layer
- self.edge_layers = nn.ModuleList()
- for i in range(num_edge_layer):
- self.edge_layers.append(AGNNConv())
-
- self.node_layers = nn.ModuleList()
- for i in range(num_node_layer):
- self.node_layers.append(GraphConv(in_feats=dim, out_feats=dim, dropout=dropout, activation=th.tanh))
- self.linear = nn.Linear(in_features=dim, out_features=out_dim, bias=False)
- self.dropout = nn.Dropout(dropout)
- self.init_para()
-
- def init_para(self):
- return
-
- def forward(self, hg, n_feats, *args, **kwargs):
- r"""
- First, apply edge_layer in cl_graph to get edge embedding.
- Then, propagate node and edge features through GraphConv.
- """
- # For full graph training, directly use the graph
- # Forward of n layers of CompGraphConv
- h = self.cl_graph.ndata['h']
- h_e = self.cl_graph.edata['w']
- for layer in self.edge_layers:
- h = th.relu(layer(self.cl_graph, h, h_e))
- h = self.dropout(h)
-
- h = self.linear_e1(h)
- edge_weight = {}
- for i, e in enumerate(hg.canonical_etypes):
- edge_weight[e] = h[i].expand(hg.num_edges(e), -1)
- if hasattr(hg, 'ntypes'):
- # edge_weight = F.embedding(hg.edata[dgl.ETYPE].long(), h)
-
- # full graph training
- for layer in self.node_layers:
- n_feats = layer(hg, n_feats, edge_weight)
- else:
- # minibatch training
- pass
- for n in n_feats:
- #n_feats[n] = self.dropout(self.linear(n_feats[n]))
- n_feats[n] = self.linear(n_feats[n])
- return n_feats
-
-
- class AGNNConv(nn.Module):
- def __init__(self,
- eps=0.,
- train_eps=False,
- learn_beta=True):
- super(AGNNConv, self).__init__()
- self.initial_eps = eps
- if learn_beta:
- self.beta = nn.Parameter(th.Tensor(1))
- else:
- self.register_buffer('beta', th.Tensor(1))
- self.learn_beta = learn_beta
- if train_eps:
- self.eps = th.nn.Parameter(th.ones([eps]))
- else:
- self.register_buffer('eps', th.Tensor([eps]))
-
- self.reset_parameters()
-
- def reset_parameters(self):
- self.eps.data.fill_(self.initial_eps)
- if self.learn_beta:
- self.beta.data.fill_(1)
-
- def forward(self, graph, feat, edge_weight):
-
- with graph.local_scope():
- feat_src, feat_dst = expand_as_pair(feat, graph)
-
- graph.srcdata['norm_h'] = F.normalize(feat_src, p=2, dim=-1)
-
- e = self.beta * edge_weight
- #graph.edata['p'] = e
- graph.edata['p'] = edge_softmax(graph, e, norm_by='src')
- graph.update_all(fn.u_mul_e('norm_h', 'p', 'm'), fn.sum('m', 'h'))
- rst = graph.dstdata.pop('h')
- rst = (1 + self.eps) * feat + rst
- return rst
-
-
- class GraphConv(nn.Module):
- def __init__(self,
- in_feats,
- out_feats, dropout,
- activation=None,
- ):
- super(GraphConv, self).__init__()
- self._in_feats = in_feats
- self._out_feats = out_feats
-
- self.weight1 = nn.Parameter(th.Tensor(in_feats, out_feats))
- #self.weight2 = nn.Parameter(th.Tensor(in_feats, out_feats))
-
- self.reset_parameters()
- self.dropout = nn.Dropout(dropout)
- self.activation = activation
-
- def reset_parameters(self):
- nn.init.xavier_uniform_(self.weight1)
- #nn.init.xavier_uniform_(self.weight2)
-
- def forward(self, hg, feat, edge_weight=None):
-
- with hg.local_scope():
- outputs = {}
- norm = {}
- aggregate_fn = fn.copy_u('h', 'm')
- if edge_weight is not None:
- #assert edge_weight.shape[0] == graph.number_of_edges()
- hg.edata['_edge_weight'] = edge_weight
- aggregate_fn = fn.u_mul_e('h', '_edge_weight', 'm')
- for e in hg.canonical_etypes:
- if e[0] == e[1]:
- hg = dgl.remove_self_loop(hg, etype=e)
- feat_src, feat_dst = expand_as_pair(feat, hg)
-
- # aggregate first then mult W
- hg.srcdata['h'] = feat_src
- for e in hg.canonical_etypes:
- stype, etype, dtype = e
- sub_graph = hg[stype, etype, dtype]
- sub_graph.update_all(aggregate_fn, fn.sum(msg='m', out='out'))
- temp = hg.ndata['out'].pop(dtype)
- degs = sub_graph.in_degrees().float().clamp(min=1)
- if isinstance(temp, dict):
- temp = temp[dtype]
- if outputs.get(dtype) is None:
- outputs[dtype] = temp
- norm[dtype] = degs
- else:
- outputs[dtype].add_(temp)
- norm[dtype].add_(degs)
-
- def _apply(ntype, h, norm):
- h = th.matmul(h+feat[ntype], self.weight1)
-
- if self.activation:
- h = self.activation(h)
- return self.dropout(h)
-
- return {ntype: _apply(ntype, h, norm) for ntype, h in outputs.items()}
-
-
|