wehub-resource-sync
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
from collections.abc import Mapping
|
|
|
|
import dgl
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
def _create_homogeneous():
|
|
s = torch.randint(0, 200, (1000,))
|
|
d = torch.randint(0, 200, (1000,))
|
|
g = dgl.graph((s, d), num_nodes=200)
|
|
reverse_eids = torch.cat([torch.arange(1000, 2000), torch.arange(0, 1000)])
|
|
seed_edges = torch.arange(0, 1000)
|
|
return g, reverse_eids, seed_edges
|
|
|
|
|
|
def _find_edges_to_exclude(g, pair_eids, degree_threshold):
|
|
src, dst = g.find_edges(pair_eids)
|
|
head_degree = g.in_degrees(src)
|
|
tail_degree = g.in_degrees(dst)
|
|
degree = torch.min(head_degree, tail_degree)
|
|
degree_mask = degree < degree_threshold
|
|
low_degree_pair_eids = pair_eids[degree_mask]
|
|
low_degree_pair_eids = torch.cat(
|
|
[low_degree_pair_eids, low_degree_pair_eids + 1000]
|
|
)
|
|
return low_degree_pair_eids
|
|
|
|
|
|
@pytest.mark.parametrize("degree_threshold", [1, 2, 3, 4, 5])
|
|
@pytest.mark.parametrize("batch_size", [1, 10, 50])
|
|
def test_spot_target_excludes(degree_threshold, batch_size):
|
|
g, reverse_eids, seed_edges = _create_homogeneous()
|
|
sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
|
|
low_degree_excluder = dgl.dataloading.SpotTarget(
|
|
g,
|
|
exclude="reverse_id",
|
|
degree_threshold=degree_threshold,
|
|
reverse_eids=reverse_eids,
|
|
)
|
|
sampler = dgl.dataloading.as_edge_prediction_sampler(
|
|
sampler,
|
|
exclude=low_degree_excluder,
|
|
negative_sampler=dgl.dataloading.negative_sampler.Uniform(1),
|
|
)
|
|
dataloader = dgl.dataloading.DataLoader(
|
|
g, seed_edges, sampler, batch_size=batch_size
|
|
)
|
|
|
|
for i, (input_nodes, pair_graph, neg_pair_graph, blocks) in enumerate(
|
|
dataloader
|
|
):
|
|
if isinstance(blocks, list):
|
|
subg = blocks[0]
|
|
else:
|
|
subg = blocks
|
|
pair_eids = pair_graph.edata[dgl.EID]
|
|
block_eids = subg.edata[dgl.EID]
|
|
edges_to_exclude = _find_edges_to_exclude(
|
|
g, pair_eids, degree_threshold
|
|
)
|
|
if edges_to_exclude is None:
|
|
continue
|
|
edges_to_exclude = dgl.utils.recursive_apply(
|
|
edges_to_exclude, lambda x: x.cpu().numpy()
|
|
)
|
|
block_eids = dgl.utils.recursive_apply(
|
|
block_eids, lambda x: x.cpu().numpy()
|
|
)
|
|
|
|
if isinstance(edges_to_exclude, Mapping):
|
|
for k in edges_to_exclude.keys():
|
|
assert not np.isin(edges_to_exclude[k], block_eids[k]).any()
|
|
else:
|
|
assert not np.isin(edges_to_exclude, block_eids).any()
|
|
|
|
if i == 10:
|
|
break
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_spot_target_excludes(degree_threshold=2, batch_size=10)
|