181 lines
7.4 KiB
Python
181 lines
7.4 KiB
Python
|
import math
|
|||
|
|
|
|||
|
|
import dgl.function as fn
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from dgl.nn import HeteroGraphConv
|
|||
|
|
from dgl.ops import edge_softmax
|
|||
|
|
from dgl.utils import expand_as_pair
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGTAttention(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, out_dim, num_heads, k_linear, q_linear, v_linear, w_att, w_msg, mu):
|
|||
|
|
"""HGT注意力模块
|
|||
|
|
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param num_heads: int 注意力头数K
|
|||
|
|
:param k_linear: nn.Linear(d_in, d_out)
|
|||
|
|
:param q_linear: nn.Linear(d_in, d_out)
|
|||
|
|
:param v_linear: nn.Linear(d_in, d_out)
|
|||
|
|
:param w_att: tensor(K, d_out/K, d_out/K)
|
|||
|
|
:param w_msg: tensor(K, d_out/K, d_out/K)
|
|||
|
|
:param mu: tensor(1)
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.out_dim = out_dim
|
|||
|
|
self.num_heads = num_heads
|
|||
|
|
self.d_k = out_dim // num_heads
|
|||
|
|
self.k_linear = k_linear
|
|||
|
|
self.q_linear = q_linear
|
|||
|
|
self.v_linear = v_linear
|
|||
|
|
self.w_att = w_att
|
|||
|
|
self.w_msg = w_msg
|
|||
|
|
self.mu = mu
|
|||
|
|
|
|||
|
|
def forward(self, g, feat):
|
|||
|
|
"""
|
|||
|
|
:param g: DGLGraph 二分图（只包含一种关系）
|
|||
|
|
:param feat: tensor(N_src, d_in) or (tensor(N_src, d_in), tensor(N_dst, d_in)) 输入特征
|
|||
|
|
:return: tensor(N_dst, d_out) 目标顶点该关于关系的表示
|
|||
|
|
"""
|
|||
|
|
with g.local_scope():
|
|||
|
|
feat_src, feat_dst = expand_as_pair(feat, g)
|
|||
|
|
# (N_src, d_in) -> (N_src, d_out) -> (N_src, K, d_out/K)
|
|||
|
|
k = self.k_linear(feat_src).view(-1, self.num_heads, self.d_k)
|
|||
|
|
v = self.v_linear(feat_src).view(-1, self.num_heads, self.d_k)
|
|||
|
|
q = self.q_linear(feat_dst).view(-1, self.num_heads, self.d_k)
|
|||
|
|
|
|||
|
|
# k[:, h] @= w_att[h] => k[n, h, j] = ∑(i) k[n, h, i] * w_att[h, i, j]
|
|||
|
|
k = torch.einsum('nhi,hij->nhj', k, self.w_att)
|
|||
|
|
v = torch.einsum('nhi,hij->nhj', v, self.w_msg)
|
|||
|
|
|
|||
|
|
g.srcdata.update({'k': k, 'v': v})
|
|||
|
|
g.dstdata['q'] = q
|
|||
|
|
g.apply_edges(fn.v_dot_u('q', 'k', 't')) # g.edata['t']: (E, K, 1)
|
|||
|
|
attn = g.edata.pop('t').squeeze(dim=-1) * self.mu / math.sqrt(self.d_k)
|
|||
|
|
attn = edge_softmax(g, attn) # (E, K)
|
|||
|
|
self.attn = attn.detach()
|
|||
|
|
g.edata['t'] = attn.unsqueeze(dim=-1) # (E, K, 1)
|
|||
|
|
|
|||
|
|
g.update_all(fn.u_mul_e('v', 't', 'm'), fn.sum('m', 'h'))
|
|||
|
|
out = g.dstdata['h'].view(-1, self.out_dim) # (N_dst, d_out)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGTLayer(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, in_dim, out_dim, num_heads, ntypes, etypes, dropout=0.2, use_norm=True):
|
|||
|
|
"""HGT层
|
|||
|
|
|
|||
|
|
:param in_dim: int 输入特征维数
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param num_heads: int 注意力头数K
|
|||
|
|
:param ntypes: List[str] 顶点类型列表
|
|||
|
|
:param etypes: List[(str, str, str)] 规范边类型列表
|
|||
|
|
:param dropout: dropout: float, optional Dropout概率，默认为0.2
|
|||
|
|
:param use_norm: bool, optional 是否使用层归一化，默认为True
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
d_k = out_dim // num_heads
|
|||
|
|
k_linear = {ntype: nn.Linear(in_dim, out_dim) for ntype in ntypes}
|
|||
|
|
q_linear = {ntype: nn.Linear(in_dim, out_dim) for ntype in ntypes}
|
|||
|
|
v_linear = {ntype: nn.Linear(in_dim, out_dim) for ntype in ntypes}
|
|||
|
|
w_att = {r[1]: nn.Parameter(torch.Tensor(num_heads, d_k, d_k)) for r in etypes}
|
|||
|
|
w_msg = {r[1]: nn.Parameter(torch.Tensor(num_heads, d_k, d_k)) for r in etypes}
|
|||
|
|
mu = {r[1]: nn.Parameter(torch.ones(num_heads)) for r in etypes}
|
|||
|
|
self.reset_parameters(w_att, w_msg)
|
|||
|
|
self.conv = HeteroGraphConv({
|
|||
|
|
etype: HGTAttention(
|
|||
|
|
out_dim, num_heads, k_linear[stype], q_linear[dtype], v_linear[stype],
|
|||
|
|
w_att[etype], w_msg[etype], mu[etype]
|
|||
|
|
) for stype, etype, dtype in etypes
|
|||
|
|
}, 'mean')
|
|||
|
|
|
|||
|
|
self.a_linear = nn.ModuleDict({ntype: nn.Linear(out_dim, out_dim) for ntype in ntypes})
|
|||
|
|
self.skip = nn.ParameterDict({ntype: nn.Parameter(torch.ones(1)) for ntype in ntypes})
|
|||
|
|
self.drop = nn.Dropout(dropout)
|
|||
|
|
|
|||
|
|
self.use_norm = use_norm
|
|||
|
|
if use_norm:
|
|||
|
|
self.norms = nn.ModuleDict({ntype: nn.LayerNorm(out_dim) for ntype in ntypes})
|
|||
|
|
|
|||
|
|
def reset_parameters(self, w_att, w_msg):
|
|||
|
|
for etype in w_att:
|
|||
|
|
nn.init.xavier_uniform_(w_att[etype])
|
|||
|
|
nn.init.xavier_uniform_(w_msg[etype])
|
|||
|
|
|
|||
|
|
def forward(self, g, feats):
|
|||
|
|
"""
|
|||
|
|
:param g: DGLGraph 异构图
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in)] 顶点类型到输入顶点特征的映射
|
|||
|
|
:return: Dict[str, tensor(N_i, d_out)] 顶点类型到输出特征的映射
|
|||
|
|
"""
|
|||
|
|
if g.is_block:
|
|||
|
|
feats_dst = {ntype: feats[ntype][:g.num_dst_nodes(ntype)] for ntype in feats}
|
|||
|
|
else:
|
|||
|
|
feats_dst = feats
|
|||
|
|
with g.local_scope():
|
|||
|
|
# 第1步：异构互注意力+异构消息传递+目标相关的聚集
|
|||
|
|
hs = self.conv(g, feats) # {ntype: tensor(N_i, d_out)}
|
|||
|
|
|
|||
|
|
# 第2步：残差连接
|
|||
|
|
out_feats = {}
|
|||
|
|
for ntype in g.dsttypes:
|
|||
|
|
if g.num_dst_nodes(ntype) == 0:
|
|||
|
|
continue
|
|||
|
|
alpha = torch.sigmoid(self.skip[ntype])
|
|||
|
|
trans_out = self.drop(self.a_linear[ntype](hs[ntype]))
|
|||
|
|
out = alpha * trans_out + (1 - alpha) * feats_dst[ntype]
|
|||
|
|
out_feats[ntype] = self.norms[ntype](out) if self.use_norm else out
|
|||
|
|
return out_feats
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGT(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self, in_dims, hidden_dim, out_dim, num_heads, ntypes, etypes,
|
|||
|
|
predict_ntype, num_layers, dropout=0.2, use_norm=True):
|
|||
|
|
"""HGT模型
|
|||
|
|
|
|||
|
|
:param in_dims: Dict[str, int] 顶点类型到输入特征维数的映射
|
|||
|
|
:param hidden_dim: int 隐含特征维数
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param num_heads: int 注意力头数K
|
|||
|
|
:param ntypes: List[str] 顶点类型列表
|
|||
|
|
:param etypes: List[(str, str, str)] 规范边类型列表
|
|||
|
|
:param predict_ntype: str 待预测顶点类型
|
|||
|
|
:param num_layers: int 层数
|
|||
|
|
:param dropout: dropout: float, optional Dropout概率，默认为0.2
|
|||
|
|
:param use_norm: bool, optional 是否使用层归一化，默认为True
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.predict_ntype = predict_ntype
|
|||
|
|
self.adapt_fcs = nn.ModuleDict({
|
|||
|
|
ntype: nn.Linear(in_dim, hidden_dim) for ntype, in_dim in in_dims.items()
|
|||
|
|
})
|
|||
|
|
self.layers = nn.ModuleList([
|
|||
|
|
HGTLayer(hidden_dim, hidden_dim, num_heads, ntypes, etypes, dropout, use_norm)
|
|||
|
|
for _ in range(num_layers)
|
|||
|
|
])
|
|||
|
|
self.predict = nn.Linear(hidden_dim, out_dim)
|
|||
|
|
|
|||
|
|
def forward(self, blocks, feats):
|
|||
|
|
"""
|
|||
|
|
:param blocks: List[DGLBlock]
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in)] 顶点类型到输入顶点特征的映射
|
|||
|
|
:return: tensor(N_i, d_out) 待预测顶点的最终嵌入
|
|||
|
|
"""
|
|||
|
|
hs = {ntype: F.gelu(self.adapt_fcs[ntype](feats[ntype])) for ntype in feats}
|
|||
|
|
for i in range(len(self.layers)):
|
|||
|
|
hs = self.layers[i](blocks[i], hs) # {ntype: tensor(N_i, d_hid)}
|
|||
|
|
out = self.predict(hs[self.predict_ntype]) # tensor(N_i, d_out)
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGTFull(HGT):
|
|||
|
|
|
|||
|
|
def forward(self, g, feats):
|
|||
|
|
return super().forward([g] * len(self.layers), feats)
|