127 lines
5.3 KiB
Python
127 lines
5.3 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
from dgl.dataloading import MultiLayerFullNeighborSampler, NodeDataLoader
|
||
|
||
from ..heco.model import PositiveGraphEncoder, Contrast
|
||
from ..rhgnn.model import RHGNN
|
||
|
||
|
||
class RHCO(nn.Module):
|
||
|
||
def __init__(
|
||
self, in_dims, hidden_dim, out_dim, rel_hidden_dim, num_heads,
|
||
ntypes, etypes, predict_ntype, num_layers, dropout, num_pos_graphs, tau, lambda_):
|
||
"""基于对比学习的关系感知异构图神经网络RHCO
|
||
|
||
:param in_dims: Dict[str, int] 顶点类型到输入特征维数的映射
|
||
:param hidden_dim: int 隐含特征维数
|
||
:param out_dim: int 输出特征维数
|
||
:param rel_hidden_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 输入特征dropout
|
||
:param num_pos_graphs: int 正样本图个数M
|
||
:param tau: float 温度参数τ
|
||
:param lambda_: float 0~1之间，网络结构视图损失的系数λ（元路径视图损失的系数为1-λ）
|
||
"""
|
||
super().__init__()
|
||
self.hidden_dim = hidden_dim
|
||
self.predict_ntype = predict_ntype
|
||
self.sc_encoder = RHGNN(
|
||
in_dims, hidden_dim, hidden_dim, rel_hidden_dim, rel_hidden_dim, num_heads,
|
||
ntypes, etypes, predict_ntype, num_layers, dropout
|
||
)
|
||
self.pg_encoder = PositiveGraphEncoder(
|
||
num_pos_graphs, in_dims[predict_ntype], hidden_dim, dropout
|
||
)
|
||
self.contrast = Contrast(hidden_dim, tau, lambda_)
|
||
self.predict = nn.Linear(hidden_dim, out_dim)
|
||
self.reset_parameters()
|
||
|
||
def reset_parameters(self):
|
||
gain = nn.init.calculate_gain('relu')
|
||
nn.init.xavier_normal_(self.predict.weight, gain)
|
||
|
||
def forward(self, blocks, feats, mgs, mg_feats, pos):
|
||
"""
|
||
:param blocks: List[DGLBlock]
|
||
:param feats: Dict[str, tensor(N_i, d_in)] 顶点类型到输入特征的映射
|
||
:param mgs: List[DGLBlock] 正样本图，len(mgs)=元路径数量=目标顶点邻居类型数S≠模型层数
|
||
:param mg_feats: List[tensor(N_pos_src, d_in)] 正样本图源顶点的输入特征
|
||
:param pos: tensor(B, N) 布尔张量，每个顶点的正样本
|
||
（B是batch大小，真正的目标顶点；N是B个目标顶点加上其正样本后的顶点数）
|
||
:return: float, tensor(B, d_out) 对比损失，目标顶点输出特征
|
||
"""
|
||
z_sc = self.sc_encoder(blocks, feats) # (N, d_hid)
|
||
z_pg = self.pg_encoder(mgs, mg_feats) # (N, d_hid)
|
||
loss = self.contrast(z_sc, z_pg, pos)
|
||
return loss, self.predict(z_sc[:pos.shape[0]])
|
||
|
||
@torch.no_grad()
|
||
def get_embeds(self, g, batch_size, device):
|
||
"""计算目标顶点的最终嵌入(z_sc)
|
||
|
||
:param g: DGLGraph 异构图
|
||
:param batch_size: int 批大小
|
||
:param device torch.device GPU设备
|
||
:return: tensor(N_tgt, d_out) 目标顶点的最终嵌入
|
||
"""
|
||
sampler = MultiLayerFullNeighborSampler(len(self.sc_encoder.layers))
|
||
loader = NodeDataLoader(
|
||
g, {self.predict_ntype: g.nodes(self.predict_ntype)}, sampler,
|
||
device=device, batch_size=batch_size
|
||
)
|
||
embeds = torch.zeros(g.num_nodes(self.predict_ntype), self.hidden_dim, device=device)
|
||
for input_nodes, output_nodes, blocks in loader:
|
||
z_sc = self.sc_encoder(blocks, blocks[0].srcdata['feat'])
|
||
embeds[output_nodes[self.predict_ntype]] = z_sc
|
||
return self.predict(embeds)
|
||
|
||
|
||
class RHCOFull(RHCO):
|
||
"""Full-batch RHCO"""
|
||
|
||
def forward(self, g, feats, mgs, mg_feat, pos):
|
||
return super().forward(
|
||
[g] * len(self.sc_encoder.layers), feats, mgs, [mg_feat] * len(mgs), pos
|
||
)
|
||
|
||
@torch.no_grad()
|
||
def get_embeds(self, g, *args):
|
||
return self.predict(self.sc_encoder([g] * len(self.sc_encoder.layers), g.ndata['feat']))
|
||
|
||
|
||
class RHCOsc(RHCO):
|
||
"""RHCO消融实验变体：仅使用网络结构编码器"""
|
||
|
||
def forward(self, blocks, feats, mgs, mg_feats, pos):
|
||
z_sc = self.sc_encoder(blocks, feats) # (N, d_hid)
|
||
loss = self.contrast(z_sc, z_sc, pos)
|
||
return loss, self.predict(z_sc[:pos.shape[0]])
|
||
|
||
|
||
class RHCOpg(RHCO):
|
||
"""RHCO消融实验变体：仅使用正样本图编码器"""
|
||
|
||
def forward(self, blocks, feats, mgs, mg_feats, pos):
|
||
z_pg = self.pg_encoder(mgs, mg_feats) # (N, d_hid)
|
||
loss = self.contrast(z_pg, z_pg, pos)
|
||
return loss, self.predict(z_pg[:pos.shape[0]])
|
||
|
||
def get_embeds(self, mgs, feat, batch_size, device):
|
||
sampler = MultiLayerFullNeighborSampler(1)
|
||
mg_loaders = [
|
||
NodeDataLoader(mg, mg.nodes(self.predict_ntype), sampler, device=device, batch_size=batch_size)
|
||
for mg in mgs
|
||
]
|
||
embeds = torch.zeros(mgs[0].num_nodes(self.predict_ntype), self.hidden_dim, device=device)
|
||
for mg_blocks in zip(*mg_loaders):
|
||
output_nodes = mg_blocks[0][1]
|
||
mg_feats = [feat[i] for i, _, _ in mg_blocks]
|
||
mg_blocks = [b[0] for _, _, b in mg_blocks]
|
||
embeds[output_nodes] = self.pg_encoder(mg_blocks, mg_feats)
|
||
return self.predict(embeds)
|