29 lines
1.3 KiB
Python
29 lines
1.3 KiB
Python
import torch
|
||
from dgl.dataloading import MultiLayerNeighborSampler
|
||
|
||
|
||
class PositiveSampler(MultiLayerNeighborSampler):
|
||
|
||
def __init__(self, fanouts, pos):
|
||
"""用于HeCo模型的邻居采样器
|
||
|
||
对于每个batch的目标顶点，将其正样本添加到目标顶点并生成block
|
||
|
||
:param fanouts: 每层的邻居采样数（见MultiLayerNeighborSampler）
|
||
:param pos: tensor(N, T_pos) 每个顶点的正样本id，N是目标顶点数
|
||
"""
|
||
super().__init__(fanouts)
|
||
self.pos = pos
|
||
|
||
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
|
||
# 如果g是异构图则seed_nodes是字典，应当只有目标顶点类型
|
||
if not g.is_homogeneous:
|
||
assert len(seed_nodes) == 1, 'PositiveSampler: 异构图只能指定目标顶点这一种类型'
|
||
ntype, seed_nodes = next(iter(seed_nodes.items()))
|
||
pos_samples = self.pos[seed_nodes].flatten() # (B, T_pos) -> (B*T_pos,)
|
||
added = list(set(pos_samples.tolist()) - set(seed_nodes.tolist()))
|
||
seed_nodes = torch.cat([seed_nodes, torch.tensor(added, device=seed_nodes.device)])
|
||
if not g.is_homogeneous:
|
||
seed_nodes = {ntype: seed_nodes}
|
||
return super().sample_blocks(g, seed_nodes, exclude_eids)
|