96 lines
3.8 KiB
Python
96 lines
3.8 KiB
Python
|
import torch.nn as nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from dgl.nn import HeteroGraphConv, GraphConv
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RelGraphConv(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(self, in_dim, out_dim, ntypes, etypes, activation=None, dropout=0.0):
|
|||
|
|
"""R-GCN层（用于异构图）
|
|||
|
|
|
|||
|
|
:param in_dim: 输入特征维数
|
|||
|
|
:param out_dim: 输出特征维数
|
|||
|
|
:param ntypes: List[str] 顶点类型列表
|
|||
|
|
:param etypes: List[str] 边类型列表
|
|||
|
|
:param activation: callable, optional 激活函数，默认为None
|
|||
|
|
:param dropout: float, optional Dropout概率，默认为0
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.activation = activation
|
|||
|
|
self.dropout = nn.Dropout(dropout)
|
|||
|
|
|
|||
|
|
self.conv = HeteroGraphConv({
|
|||
|
|
etype: GraphConv(in_dim, out_dim, norm='right', bias=False)
|
|||
|
|
for etype in etypes
|
|||
|
|
}, 'sum')
|
|||
|
|
self.loop_weight = nn.ModuleDict({
|
|||
|
|
ntype: nn.Linear(in_dim, out_dim, bias=False) for ntype in ntypes
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
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: feat[:g.num_dst_nodes(ntype)] for ntype, feat in feats.items()}
|
|||
|
|
else:
|
|||
|
|
feats_dst = feats
|
|||
|
|
out = self.conv(g, (feats, feats_dst)) # Dict[ntype, (N_i, d_out)]
|
|||
|
|
for ntype in out:
|
|||
|
|
out[ntype] += self.loop_weight[ntype](feats_dst[ntype])
|
|||
|
|
if self.activation:
|
|||
|
|
out[ntype] = self.activation(out[ntype])
|
|||
|
|
out[ntype] = self.dropout(out[ntype])
|
|||
|
|
return out
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RGCN(nn.Module):
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self, in_dim, hidden_dim, out_dim, input_ntypes, num_nodes, etypes, predict_ntype,
|
|||
|
|
num_layers=2, dropout=0.0):
|
|||
|
|
"""R-GCN模型
|
|||
|
|
|
|||
|
|
:param in_dim: int 输入特征维数
|
|||
|
|
:param hidden_dim: int 隐含特征维数
|
|||
|
|
:param out_dim: int 输出特征维数
|
|||
|
|
:param input_ntypes: List[str] 有输入特征的顶点类型列表
|
|||
|
|
:param num_nodes: Dict[str, int] 顶点类型到顶点数的映射
|
|||
|
|
:param etypes: List[str] 边类型列表
|
|||
|
|
:param predict_ntype: str 待预测顶点类型
|
|||
|
|
:param num_layers: int, optional 层数，默认为2
|
|||
|
|
:param dropout: float, optional Dropout概率，默认为0
|
|||
|
|
"""
|
|||
|
|
super().__init__()
|
|||
|
|
self.embeds = nn.ModuleDict({
|
|||
|
|
ntype: nn.Embedding(num_nodes[ntype], in_dim)
|
|||
|
|
for ntype in num_nodes if ntype not in input_ntypes
|
|||
|
|
})
|
|||
|
|
ntypes = list(num_nodes)
|
|||
|
|
self.layers = nn.ModuleList()
|
|||
|
|
self.layers.append(RelGraphConv(in_dim, hidden_dim, ntypes, etypes, F.relu, dropout))
|
|||
|
|
for i in range(num_layers - 2):
|
|||
|
|
self.layers.append(RelGraphConv(hidden_dim, hidden_dim, ntypes, etypes, F.relu, dropout))
|
|||
|
|
self.layers.append(RelGraphConv(hidden_dim, out_dim, ntypes, etypes))
|
|||
|
|
self.predict_ntype = predict_ntype
|
|||
|
|
self.reset_parameters()
|
|||
|
|
|
|||
|
|
def reset_parameters(self):
|
|||
|
|
gain = nn.init.calculate_gain('relu')
|
|||
|
|
for k in self.embeds:
|
|||
|
|
nn.init.xavier_uniform_(self.embeds[k].weight, gain=gain)
|
|||
|
|
|
|||
|
|
def forward(self, g, feats):
|
|||
|
|
"""
|
|||
|
|
:param g: DGLGraph 异构图
|
|||
|
|
:param feats: Dict[str, tensor(N_i, d_in_i)] （部分）顶点类型到输入特征的映射
|
|||
|
|
:return: Dict[str, tensor(N_i, d_out)] 顶点类型到顶点嵌入的映射
|
|||
|
|
"""
|
|||
|
|
for k in self.embeds:
|
|||
|
|
feats[k] = self.embeds[k].weight
|
|||
|
|
for i in range(len(self.layers)):
|
|||
|
|
feats = self.layers[i](g, feats) # Dict[ntype, (N_i, d_hid)]
|
|||
|
|
return feats[self.predict_ntype]
|