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)
|