292 lines
13 KiB
Python
292 lines
13 KiB
Python
|
import dgl.function as fn
|
|||
|
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from dgl.dataloading import MultiLayerFullNeighborSampler, NodeDataLoader
|
|||
|
|
from dgl.ops import edge_softmax
|
|||
|
|
from dgl.utils import expand_as_pair
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MicroConv(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self, out_dim, num_heads, fc_src, fc_dst, attn_src,
|
|||
|
|
feat_drop=0.0, negative_slope=0.2, activation=None):
|
|||
|
|
"""微观层次卷积
|
|||
|
|
|
|||
|
|
针对一种关系（边类型）R=<stype, etype, dtype>，聚集关系R下的邻居信息，得到关系R关于dtype类型顶点的表示
|
|||
|
|
（特征转换矩阵和注意力向量是与顶点类型相关的，除此之外与GAT完全相同）
|
|||
|
|
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param num_heads: int 注意力头数K
|
|||
|
|
:param fc_src: nn.Linear(d_in, K*d_out) 源顶点特征转换模块
|
|||
|
|
:param fc_dst: nn.Linear(d_in, K*d_out) 目标顶点特征转换模块
|
|||
|
|
:param attn_src: nn.Parameter(K, 2d_out) 源顶点类型对应的注意力向量
|
|||
|
|
:param feat_drop: float, optional 输入特征Dropout概率，默认为0
|
|||
|
|
:param negative_slope: float, optional LeakyReLU负斜率，默认为0.2
|
|||
|
|
:param activation: callable, optional 用于输出特征的激活函数，默认为None
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.out_dim = out_dim
|
|||
|
|
self.num_heads = num_heads
|
|||
|
|
self.fc_src = fc_src
|
|||
|
|
self.fc_dst = fc_dst
|
|||
|
|
self.attn_src = attn_src
|
|||
|
|
self.feat_drop = nn.Dropout(feat_drop)
|
|||
|
|
self.leaky_relu = nn.LeakyReLU(negative_slope)
|
|||
|
|
self.activation = activation
|
|||
|
|
|
|||
|
|
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, K*d_out) 该关系关于目标顶点的表示
|
|||
|
|
"""
|
|||
|
|
with g.local_scope():
|
|||
|
|
feat_src, feat_dst = expand_as_pair(feat, g)
|
|||
|
|
feat_src = self.fc_src(self.feat_drop(feat_src)).view(-1, self.num_heads, self.out_dim)
|
|||
|
|
feat_dst = self.fc_dst(self.feat_drop(feat_dst)).view(-1, self.num_heads, self.out_dim)
|
|||
|
|
|
|||
|
|
# a^T (z_u || z_v) = (a_l^T || a_r^T) (z_u || z_v) = a_l^T z_u + a_r^T z_v = el + er
|
|||
|
|
el = (feat_src * self.attn_src[:, :self.out_dim]).sum(dim=-1, keepdim=True) # (N_src, K, 1)
|
|||
|
|
er = (feat_dst * self.attn_src[:, self.out_dim:]).sum(dim=-1, keepdim=True) # (N_dst, K, 1)
|
|||
|
|
g.srcdata.update({'ft': feat_src, 'el': el})
|
|||
|
|
g.dstdata['er'] = er
|
|||
|
|
g.apply_edges(fn.u_add_v('el', 'er', 'e'))
|
|||
|
|
e = self.leaky_relu(g.edata.pop('e'))
|
|||
|
|
g.edata['a'] = edge_softmax(g, e) # (E, K, 1)
|
|||
|
|
|
|||
|
|
# 消息传递
|
|||
|
|
g.update_all(fn.u_mul_e('ft', 'a', 'm'), fn.sum('m', 'ft'))
|
|||
|
|
ret = g.dstdata['ft'].view(-1, self.num_heads * self.out_dim)
|
|||
|
|
if self.activation:
|
|||
|
|
ret = self.activation(ret)
|
|||
|
|
return ret
|
|||
|
|
|
|||
|
|
|
|||
|
|
class MacroConv(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, out_dim, num_heads, fc_node, fc_rel, attn, dropout=0.0, negative_slope=0.2):
|
|||
|
|
"""宏观层次卷积
|
|||
|
|
|
|||
|
|
针对所有关系（边类型），将每种类型的顶点关联的所有关系关于该类型顶点的表示组合起来
|
|||
|
|
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param num_heads: int 注意力头数K
|
|||
|
|
:param fc_node: Dict[str, nn.Linear(d_in, K*d_out)] 顶点类型到顶点特征转换模块的映射
|
|||
|
|
:param fc_rel: Dict[str, nn.Linear(K*d_out, K*d_out)] 关系到关系表示转换模块的映射
|
|||
|
|
:param attn: nn.Parameter(K, 2d_out)
|
|||
|
|
:param dropout: float, optional Dropout概率，默认为0
|
|||
|
|
:param negative_slope: float, optional LeakyReLU负斜率，默认为0.2
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.out_dim = out_dim
|
|||
|
|
self.num_heads = num_heads
|
|||
|
|
self.fc_node = fc_node
|
|||
|
|
self.fc_rel = fc_rel
|
|||
|
|
self.attn = attn
|
|||
|
|
self.dropout = nn.Dropout(dropout)
|
|||
|
|
self.leaky_relu = nn.LeakyReLU(negative_slope)
|
|||
|
|
|
|||
|
|
def forward(self, node_feats, rel_feats):
|
|||
|
|
"""
|
|||
|
|
:param node_feats: Dict[str, tensor(N_i, d_in) 顶点类型到输入顶点特征的映射
|
|||
|
|
:param rel_feats: Dict[(str, str, str), tensor(N_i, K*d_out)]
|
|||
|
|
关系(stype, etype, dtype)到关系关于其终点类型的表示的映射
|
|||
|
|
:return: Dict[str, tensor(N_i, K*d_out)] 顶点类型到最终顶点嵌入的映射
|
|||
|
|
"""
|
|||
|
|
node_feats = {
|
|||
|
|
ntype: self.fc_node[ntype](feat).view(-1, self.num_heads, self.out_dim)
|
|||
|
|
for ntype, feat in node_feats.items()
|
|||
|
|
}
|
|||
|
|
rel_feats = {
|
|||
|
|
r: self.fc_rel[r[1]](feat).view(-1, self.num_heads, self.out_dim)
|
|||
|
|
for r, feat in rel_feats.items()
|
|||
|
|
}
|
|||
|
|
out_feats = {}
|
|||
|
|
for ntype, node_feat in node_feats.items():
|
|||
|
|
rel_node_feats = [feat for rel, feat in rel_feats.items() if rel[2] == ntype]
|
|||
|
|
if not rel_node_feats:
|
|||
|
|
continue
|
|||
|
|
elif len(rel_node_feats) == 1:
|
|||
|
|
out_feats[ntype] = rel_node_feats[0].view(-1, self.num_heads * self.out_dim)
|
|||
|
|
else:
|
|||
|
|
rel_node_feats = torch.stack(rel_node_feats, dim=0) # (R, N_i, K, d_out)
|
|||
|
|
cat_feats = torch.cat(
|
|||
|
|
(node_feat.repeat(rel_node_feats.shape[0], 1, 1, 1), rel_node_feats), dim=-1
|
|||
|
|
) # (R, N_i, K, 2d_out)
|
|||
|
|
attn_scores = self.leaky_relu((self.attn * cat_feats).sum(dim=-1, keepdim=True))
|
|||
|
|
attn_scores = F.softmax(attn_scores, dim=0) # (R, N_i, K, 1)
|
|||
|
|
out_feat = (attn_scores * rel_node_feats).sum(dim=0) # (N_i, K, d_out)
|
|||
|
|
out_feats[ntype] = self.dropout(out_feat.reshape(-1, self.num_heads * self.out_dim))
|
|||
|
|
return out_feats
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGConvLayer(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, in_dim, out_dim, num_heads, ntypes, etypes, dropout=0.0, residual=True):
|
|||
|
|
"""HGConv层
|
|||
|
|
|
|||
|
|
: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: float, optional Dropout概率，默认为0
|
|||
|
|
:param residual: bool, optional 是否使用残差连接，默认True
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
# 微观层次卷积的参数
|
|||
|
|
micro_fc = {ntype: nn.Linear(in_dim, num_heads * out_dim, bias=False) for ntype in ntypes}
|
|||
|
|
micro_attn = {
|
|||
|
|
ntype: nn.Parameter(torch.FloatTensor(size=(num_heads, 2 * out_dim)))
|
|||
|
|
for ntype in ntypes
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 宏观层次卷积的参数
|
|||
|
|
macro_fc_node = nn.ModuleDict({
|
|||
|
|
ntype: nn.Linear(in_dim, num_heads * out_dim, bias=False) for ntype in ntypes
|
|||
|
|
})
|
|||
|
|
macro_fc_rel = nn.ModuleDict({
|
|||
|
|
r[1]: nn.Linear(num_heads * out_dim, num_heads * out_dim, bias=False)
|
|||
|
|
for r in etypes
|
|||
|
|
})
|
|||
|
|
macro_attn = nn.Parameter(torch.FloatTensor(size=(num_heads, 2 * out_dim)))
|
|||
|
|
|
|||
|
|
self.micro_conv = nn.ModuleDict({
|
|||
|
|
etype: MicroConv(
|
|||
|
|
out_dim, num_heads, micro_fc[stype],
|
|||
|
|
micro_fc[dtype], micro_attn[stype], dropout, activation=F.relu
|
|||
|
|
) for stype, etype, dtype in etypes
|
|||
|
|
})
|
|||
|
|
self.macro_conv = MacroConv(
|
|||
|
|
out_dim, num_heads, macro_fc_node, macro_fc_rel, macro_attn, dropout
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
self.residual = residual
|
|||
|
|
if residual:
|
|||
|
|
self.res_fc = nn.ModuleDict({
|
|||
|
|
ntype: nn.Linear(in_dim, num_heads * out_dim) for ntype in ntypes
|
|||
|
|
})
|
|||
|
|
self.res_weight = nn.ParameterDict({
|
|||
|
|
ntype: nn.Parameter(torch.rand(1)) for ntype in ntypes
|
|||
|
|
})
|
|||
|
|
self.reset_parameters(micro_fc, micro_attn, macro_fc_node, macro_fc_rel, macro_attn)
|
|||
|
|
|
|||
|
|
def reset_parameters(self, micro_fc, micro_attn, macro_fc_node, macro_fc_rel, macro_attn):
|
|||
|
|
gain = nn.init.calculate_gain('relu')
|
|||
|
|
for ntype in micro_fc:
|
|||
|
|
nn.init.xavier_normal_(micro_fc[ntype].weight, gain=gain)
|
|||
|
|
nn.init.xavier_normal_(micro_attn[ntype], gain=gain)
|
|||
|
|
nn.init.xavier_normal_(macro_fc_node[ntype].weight, gain=gain)
|
|||
|
|
if self.residual:
|
|||
|
|
nn.init.xavier_normal_(self.res_fc[ntype].weight, gain=gain)
|
|||
|
|
for etype in macro_fc_rel:
|
|||
|
|
nn.init.xavier_normal_(macro_fc_rel[etype].weight, gain=gain)
|
|||
|
|
nn.init.xavier_normal_(macro_attn, gain=gain)
|
|||
|
|
|
|||
|
|
def forward(self, g, feats):
|
|||
|
|
"""
|
|||
|
|
:param g: DGLGraph 异构图
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in)] 顶点类型到输入顶点特征的映射
|
|||
|
|
:return: Dict[str, tensor(N_i, K*d_out)] 顶点类型到最终顶点嵌入的映射
|
|||
|
|
"""
|
|||
|
|
if g.is_block:
|
|||
|
|
feats_dst = {ntype: feats[ntype][:g.num_dst_nodes(ntype)] for ntype in feats}
|
|||
|
|
else:
|
|||
|
|
feats_dst = feats
|
|||
|
|
rel_feats = {
|
|||
|
|
(stype, etype, dtype): self.micro_conv[etype](
|
|||
|
|
g[stype, etype, dtype], (feats[stype], feats_dst[dtype])
|
|||
|
|
)
|
|||
|
|
for stype, etype, dtype in g.canonical_etypes
|
|||
|
|
if g.num_edges((stype, etype, dtype)) > 0
|
|||
|
|
} # {rel: tensor(N_i, K*d_out)}
|
|||
|
|
out_feats = self.macro_conv(feats_dst, rel_feats) # {ntype: tensor(N_i, K*d_out)}
|
|||
|
|
if self.residual:
|
|||
|
|
for ntype in out_feats:
|
|||
|
|
alpha = torch.sigmoid(self.res_weight[ntype])
|
|||
|
|
inherit_feat = self.res_fc[ntype](feats_dst[ntype])
|
|||
|
|
out_feats[ntype] = alpha * out_feats[ntype] + (1 - alpha) * inherit_feat
|
|||
|
|
return out_feats
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGConv(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self, in_dims, hidden_dim, out_dim, num_heads, ntypes, etypes, predict_ntype,
|
|||
|
|
num_layers, dropout=0.0, residual=True):
|
|||
|
|
"""HGConv模型
|
|||
|
|
|
|||
|
|
: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: float, optional Dropout概率，默认为0
|
|||
|
|
:param residual: bool, optional 是否使用残差连接，默认True
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.d = num_heads * hidden_dim
|
|||
|
|
self.predict_ntype = predict_ntype
|
|||
|
|
# 对齐输入特征维数
|
|||
|
|
self.fc_in = nn.ModuleDict({
|
|||
|
|
ntype: nn.Linear(in_dim, num_heads * hidden_dim) for ntype, in_dim in in_dims.items()
|
|||
|
|
})
|
|||
|
|
self.layers = nn.ModuleList([
|
|||
|
|
HGConvLayer(
|
|||
|
|
num_heads * hidden_dim, hidden_dim, num_heads, ntypes, etypes, dropout, residual
|
|||
|
|
) for _ in range(num_layers)
|
|||
|
|
])
|
|||
|
|
self.classifier = nn.Linear(num_heads * hidden_dim, out_dim)
|
|||
|
|
|
|||
|
|
def forward(self, blocks, feats):
|
|||
|
|
"""
|
|||
|
|
:param blocks: List[DGLBlock]
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in_i)] 顶点类型到输入顶点特征的映射
|
|||
|
|
:return: tensor(N_i, d_out) 待预测顶点的最终嵌入
|
|||
|
|
"""
|
|||
|
|
feats = {ntype: self.fc_in[ntype](feat) for ntype, feat in feats.items()}
|
|||
|
|
for i in range(len(self.layers)):
|
|||
|
|
feats = self.layers[i](blocks[i], feats) # {ntype: tensor(N_i, K*d_hid)}
|
|||
|
|
return self.classifier(feats[self.predict_ntype])
|
|||
|
|
|
|||
|
|
@torch.no_grad()
|
|||
|
|
def inference(self, g, feats, device, batch_size):
|
|||
|
|
"""离线推断所有顶点的最终嵌入（不使用邻居采样）
|
|||
|
|
|
|||
|
|
:param g: DGLGraph 异构图
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in_i)] 顶点类型到输入顶点特征的映射
|
|||
|
|
:param device: torch.device
|
|||
|
|
:param batch_size: int 批大小
|
|||
|
|
:return: tensor(N_i, d_out) 待预测顶点的最终嵌入
|
|||
|
|
"""
|
|||
|
|
g.ndata['emb'] = {ntype: self.fc_in[ntype](feat) for ntype, feat in feats.items()}
|
|||
|
|
for layer in self.layers:
|
|||
|
|
embeds = {
|
|||
|
|
ntype: torch.zeros(g.num_nodes(ntype), self.d, device=device)
|
|||
|
|
for ntype in g.ntypes
|
|||
|
|
}
|
|||
|
|
sampler = MultiLayerFullNeighborSampler(1)
|
|||
|
|
loader = NodeDataLoader(
|
|||
|
|
g, {ntype: g.nodes(ntype) for ntype in g.ntypes}, sampler, device=device,
|
|||
|
|
batch_size=batch_size, shuffle=True
|
|||
|
|
)
|
|||
|
|
for input_nodes, output_nodes, blocks in tqdm(loader):
|
|||
|
|
block = blocks[0]
|
|||
|
|
h = layer(block, block.srcdata['emb'])
|
|||
|
|
for ntype in h:
|
|||
|
|
embeds[ntype][output_nodes[ntype]] = h[ntype]
|
|||
|
|
g.ndata['emb'] = embeds
|
|||
|
|
return self.classifier(g.nodes[self.predict_ntype].data['emb'])
|
|||
|
|
|
|||
|
|
|
|||
|
|
class HGConvFull(HGConv):
|
|||
|
|
|
|||
|
|
def forward(self, g, feats):
|
|||
|
|
return super().forward([g] * len(self.layers), feats)
|