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