chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1 @@
|
||||
""" DGL graphbolt API tests"""
|
||||
@@ -0,0 +1,374 @@
|
||||
import os
|
||||
|
||||
import dgl
|
||||
import dgl.graphbolt as gb
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import scipy.sparse as sp
|
||||
import torch
|
||||
|
||||
|
||||
def rand_csc_graph(N, density, bidirection_edge=False):
|
||||
adj = sp.random(N, N, density)
|
||||
if bidirection_edge:
|
||||
adj = adj + adj.T
|
||||
adj = adj.tocsc()
|
||||
|
||||
indptr = torch.LongTensor(adj.indptr)
|
||||
indices = torch.LongTensor(adj.indices)
|
||||
|
||||
graph = gb.fused_csc_sampling_graph(indptr, indices)
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def random_homo_graph(num_nodes, num_edges):
|
||||
csc_indptr = torch.randint(0, num_edges, (num_nodes + 1,))
|
||||
csc_indptr = torch.sort(csc_indptr)[0]
|
||||
csc_indptr[0] = 0
|
||||
csc_indptr[-1] = num_edges
|
||||
indices = torch.randint(0, num_nodes, (num_edges,))
|
||||
return csc_indptr, indices
|
||||
|
||||
|
||||
def get_type_to_id(num_ntypes, num_etypes):
|
||||
ntypes = {f"n{i}": i for i in range(num_ntypes)}
|
||||
etypes = {}
|
||||
count = 0
|
||||
for n1 in range(num_ntypes):
|
||||
for n2 in range(n1, num_ntypes):
|
||||
if count >= num_etypes:
|
||||
break
|
||||
etypes.update({f"n{n1}:e{count}:n{n2}": count})
|
||||
count += 1
|
||||
return ntypes, etypes
|
||||
|
||||
|
||||
def get_ntypes_and_etypes(num_nodes, num_ntypes, num_etypes):
|
||||
ntypes = {f"n{i}": num_nodes // num_ntypes for i in range(num_ntypes)}
|
||||
if num_nodes % num_ntypes != 0:
|
||||
ntypes["n0"] += num_nodes % num_ntypes
|
||||
etypes = []
|
||||
count = 0
|
||||
while count < num_etypes:
|
||||
for n1 in range(num_ntypes):
|
||||
for n2 in range(num_ntypes):
|
||||
if count >= num_etypes:
|
||||
break
|
||||
etypes.append((f"n{n1}", f"e{count}", f"n{n2}"))
|
||||
count += 1
|
||||
return ntypes, etypes
|
||||
|
||||
|
||||
def random_hetero_graph(num_nodes, num_edges, num_ntypes, num_etypes):
|
||||
ntypes, etypes = get_ntypes_and_etypes(num_nodes, num_ntypes, num_etypes)
|
||||
edges = {}
|
||||
for step, etype in enumerate(etypes):
|
||||
src_ntype, _, dst_ntype = etype
|
||||
num_e = num_edges // num_etypes + (
|
||||
0 if step != 0 else num_edges % num_etypes
|
||||
)
|
||||
if ntypes[src_ntype] == 0 or ntypes[dst_ntype] == 0:
|
||||
continue
|
||||
src = torch.randint(0, ntypes[src_ntype], (num_e,))
|
||||
dst = torch.randint(0, ntypes[dst_ntype], (num_e,))
|
||||
|
||||
edges[etype] = (src, dst)
|
||||
|
||||
gb_g = gb.from_dglgraph(dgl.heterograph(edges, ntypes))
|
||||
return (
|
||||
gb_g.csc_indptr,
|
||||
gb_g.indices,
|
||||
gb_g.node_type_offset,
|
||||
gb_g.type_per_edge,
|
||||
gb_g.node_type_to_id,
|
||||
gb_g.edge_type_to_id,
|
||||
)
|
||||
|
||||
|
||||
def random_homo_graphbolt_graph(
|
||||
test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt="csv"
|
||||
):
|
||||
"""Generate random graphbolt version homograph"""
|
||||
# Generate random edges.
|
||||
nodes = np.repeat(np.arange(num_nodes, dtype=np.int64), 5)
|
||||
neighbors = np.random.randint(
|
||||
0, num_nodes, size=(num_edges), dtype=np.int64
|
||||
)
|
||||
edges = np.stack([nodes, neighbors], axis=1)
|
||||
os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
|
||||
assert edge_fmt in [
|
||||
"numpy",
|
||||
"csv",
|
||||
], "Only numpy and csv are supported for edges."
|
||||
if edge_fmt == "csv":
|
||||
# Write into edges/edge.csv
|
||||
edges_DataFrame = pd.DataFrame(edges, columns=["src", "dst"])
|
||||
edge_path = os.path.join("edges", "edge.csv")
|
||||
edges_DataFrame.to_csv(
|
||||
os.path.join(test_dir, edge_path),
|
||||
index=False,
|
||||
header=False,
|
||||
)
|
||||
else:
|
||||
# Write into edges/edge.npy
|
||||
edges = edges.T
|
||||
edge_path = os.path.join("edges", "edge.npy")
|
||||
np.save(os.path.join(test_dir, edge_path), edges)
|
||||
|
||||
# Generate random graph edge-feats.
|
||||
edge_feats = np.random.rand(num_edges, num_classes)
|
||||
os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
|
||||
edge_feat_path = os.path.join("data", "edge-feat.npy")
|
||||
np.save(os.path.join(test_dir, edge_feat_path), edge_feats)
|
||||
|
||||
# Generate random node-feats.
|
||||
if num_classes == 1:
|
||||
node_feats = np.random.rand(num_nodes)
|
||||
else:
|
||||
node_feats = np.random.rand(num_nodes, num_classes)
|
||||
node_feat_path = os.path.join("data", "node-feat.npy")
|
||||
np.save(os.path.join(test_dir, node_feat_path), node_feats)
|
||||
|
||||
# Generate train/test/valid set.
|
||||
assert num_nodes % 4 == 0, "num_nodes must be divisible by 4"
|
||||
each_set_size = num_nodes // 4
|
||||
os.makedirs(os.path.join(test_dir, "set"), exist_ok=True)
|
||||
train_pairs = (
|
||||
np.arange(each_set_size),
|
||||
np.arange(each_set_size, 2 * each_set_size),
|
||||
)
|
||||
train_data = np.vstack(train_pairs).T.astype(edges.dtype)
|
||||
train_path = os.path.join("set", "train.npy")
|
||||
np.save(os.path.join(test_dir, train_path), train_data)
|
||||
|
||||
validation_pairs = (
|
||||
np.arange(each_set_size, 2 * each_set_size),
|
||||
np.arange(2 * each_set_size, 3 * each_set_size),
|
||||
)
|
||||
validation_data = np.vstack(validation_pairs).T.astype(edges.dtype)
|
||||
validation_path = os.path.join("set", "validation.npy")
|
||||
np.save(os.path.join(test_dir, validation_path), validation_data)
|
||||
|
||||
test_pairs = (
|
||||
np.arange(2 * each_set_size, 3 * each_set_size),
|
||||
np.arange(3 * each_set_size, 4 * each_set_size),
|
||||
)
|
||||
test_data = np.vstack(test_pairs).T.astype(edges.dtype)
|
||||
test_path = os.path.join("set", "test.npy")
|
||||
np.save(os.path.join(test_dir, test_path), test_data)
|
||||
|
||||
yaml_content = f"""
|
||||
dataset_name: {dataset_name}
|
||||
graph: # Graph structure and required attributes.
|
||||
nodes:
|
||||
- num: {num_nodes}
|
||||
edges:
|
||||
- format: {edge_fmt}
|
||||
path: {edge_path}
|
||||
feature_data:
|
||||
- domain: node
|
||||
type: null
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feat_path}
|
||||
- domain: edge
|
||||
type: null
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {edge_feat_path}
|
||||
feature_data:
|
||||
- domain: node
|
||||
type: null
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feat_path}
|
||||
- domain: edge
|
||||
type: null
|
||||
name: feat
|
||||
format: numpy
|
||||
path: {edge_feat_path}
|
||||
tasks:
|
||||
- name: link_prediction
|
||||
num_classes: {num_classes}
|
||||
train_set:
|
||||
- type: null
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {train_path}
|
||||
validation_set:
|
||||
- type: null
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {validation_path}
|
||||
test_set:
|
||||
- type: null
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {test_path}
|
||||
"""
|
||||
return yaml_content
|
||||
|
||||
|
||||
def generate_raw_data_for_hetero_dataset(
|
||||
test_dir, dataset_name, num_nodes, num_edges, num_classes, edge_fmt="csv"
|
||||
):
|
||||
# Generate edges.
|
||||
edges_path = {}
|
||||
for etype, num_edge in num_edges.items():
|
||||
src_ntype, etype_str, dst_ntype = etype
|
||||
src = torch.randint(0, num_nodes[src_ntype], (num_edge,))
|
||||
dst = torch.randint(0, num_nodes[dst_ntype], (num_edge,))
|
||||
os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
|
||||
assert edge_fmt in [
|
||||
"numpy",
|
||||
"csv",
|
||||
], "Only numpy and csv are supported for edges."
|
||||
if edge_fmt == "csv":
|
||||
# Write into edges/edge.csv
|
||||
edges = pd.DataFrame(
|
||||
np.stack([src, dst], axis=1), columns=["src", "dst"]
|
||||
)
|
||||
edge_path = os.path.join("edges", f"{etype_str}.csv")
|
||||
edges.to_csv(
|
||||
os.path.join(test_dir, edge_path),
|
||||
index=False,
|
||||
header=False,
|
||||
)
|
||||
else:
|
||||
edges = np.stack([src, dst], axis=1).T
|
||||
edge_path = os.path.join("edges", f"{etype_str}.npy")
|
||||
np.save(os.path.join(test_dir, edge_path), edges)
|
||||
edges_path[etype_str] = edge_path
|
||||
|
||||
# Generate node features.
|
||||
node_feats_path = {}
|
||||
os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
|
||||
for ntype, num_node in num_nodes.items():
|
||||
node_feat_path = os.path.join("data", f"{ntype}-feat.npy")
|
||||
node_feats = np.random.rand(num_node, num_classes)
|
||||
np.save(os.path.join(test_dir, node_feat_path), node_feats)
|
||||
node_feats_path[ntype] = node_feat_path
|
||||
|
||||
# Generate edge features.
|
||||
edge_feats_path = {}
|
||||
os.makedirs(os.path.join(test_dir, "data"), exist_ok=True)
|
||||
for etype, num_edge in num_edges.items():
|
||||
src_ntype, etype_str, dst_ntype = etype
|
||||
edge_feat_path = os.path.join("data", f"{etype_str}-feat.npy")
|
||||
edge_feats = np.random.rand(num_edge, num_classes)
|
||||
np.save(os.path.join(test_dir, edge_feat_path), edge_feats)
|
||||
edge_feats_path[etype_str] = edge_feat_path
|
||||
|
||||
# Generate train/test/valid set.
|
||||
os.makedirs(os.path.join(test_dir, "set"), exist_ok=True)
|
||||
user_ids = torch.arange(num_nodes["user"])
|
||||
np.random.shuffle(user_ids.numpy())
|
||||
num_train = int(num_nodes["user"] * 0.6)
|
||||
num_validation = int(num_nodes["user"] * 0.2)
|
||||
num_test = num_nodes["user"] - num_train - num_validation
|
||||
train_path = os.path.join("set", "train.npy")
|
||||
np.save(os.path.join(test_dir, train_path), user_ids[:num_train])
|
||||
validation_path = os.path.join("set", "validation.npy")
|
||||
np.save(
|
||||
os.path.join(test_dir, validation_path),
|
||||
user_ids[num_train : num_train + num_validation],
|
||||
)
|
||||
test_path = os.path.join("set", "test.npy")
|
||||
np.save(
|
||||
os.path.join(test_dir, test_path),
|
||||
user_ids[num_train + num_validation :],
|
||||
)
|
||||
|
||||
yaml_content = f"""
|
||||
dataset_name: {dataset_name}
|
||||
graph: # Graph structure and required attributes.
|
||||
nodes:
|
||||
- type: user
|
||||
num: {num_nodes["user"]}
|
||||
- type: item
|
||||
num: {num_nodes["item"]}
|
||||
edges:
|
||||
- type: "user:follow:user"
|
||||
format: {edge_fmt}
|
||||
path: {edges_path["follow"]}
|
||||
- type: "user:click:item"
|
||||
format: {edge_fmt}
|
||||
path: {edges_path["click"]}
|
||||
feature_data:
|
||||
- domain: node
|
||||
type: user
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feats_path["user"]}
|
||||
- domain: node
|
||||
type: item
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feats_path["item"]}
|
||||
- domain: edge
|
||||
type: "user:follow:user"
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {edge_feats_path["follow"]}
|
||||
- domain: edge
|
||||
type: "user:click:item"
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {edge_feats_path["click"]}
|
||||
feature_data:
|
||||
- domain: node
|
||||
type: user
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feats_path["user"]}
|
||||
- domain: node
|
||||
type: item
|
||||
name: feat
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {node_feats_path["item"]}
|
||||
tasks:
|
||||
- name: node_classification
|
||||
num_classes: {num_classes}
|
||||
train_set:
|
||||
- type: user
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {train_path}
|
||||
validation_set:
|
||||
- type: user
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {validation_path}
|
||||
test_set:
|
||||
- type: user
|
||||
data:
|
||||
- name: seeds
|
||||
format: numpy
|
||||
in_memory: true
|
||||
path: {test_path}
|
||||
"""
|
||||
|
||||
yaml_file = os.path.join(test_dir, "metadata.yaml")
|
||||
with open(yaml_file, "w") as f:
|
||||
f.write(yaml_content)
|
||||
@@ -0,0 +1 @@
|
||||
""" DGL graphbolt/impl tests"""
|
||||
@@ -0,0 +1,151 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def test_basic_feature_store_homo():
|
||||
a = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [4, 3]]])
|
||||
metadata = {"max_value": 3}
|
||||
|
||||
features = {}
|
||||
features[("node", None, "a")] = gb.TorchBasedFeature(a, metadata=metadata)
|
||||
features[("node", None, "b")] = gb.TorchBasedFeature(b)
|
||||
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
# Test __getitem__ to access the stored Feature.
|
||||
feature = feature_store[("node", None, "a")]
|
||||
assert isinstance(feature, gb.Feature)
|
||||
assert torch.equal(
|
||||
feature.read(),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", None, "a"),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_store.read("node", None, "b"),
|
||||
torch.tensor([[[1, 2], [3, 4]], [[2, 5], [4, 3]]]),
|
||||
)
|
||||
|
||||
# Test read with ids.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", None, "a", torch.tensor([0])),
|
||||
torch.tensor([[1, 2, 4]]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_store.read("node", None, "b", torch.tensor([0])),
|
||||
torch.tensor([[[1, 2], [3, 4]]]),
|
||||
)
|
||||
|
||||
# Test get the size and count of the entire feature.
|
||||
assert feature_store.size("node", None, "a") == torch.Size([3])
|
||||
assert feature_store.size("node", None, "b") == torch.Size([2, 2])
|
||||
assert feature_store.count("node", None, "a") == a.size(0)
|
||||
assert feature_store.count("node", None, "b") == b.size(0)
|
||||
|
||||
# Test get metadata of the feature.
|
||||
assert feature_store.metadata("node", None, "a") == metadata
|
||||
assert feature_store.metadata("node", None, "b") == {}
|
||||
|
||||
# Test __setitem__ and __contains__ of FeatureStore.
|
||||
assert ("node", None, "c") not in feature_store
|
||||
feature_store[("node", None, "c")] = feature_store[("node", None, "a")]
|
||||
assert ("node", None, "c") in feature_store
|
||||
|
||||
# Test get keys of the features.
|
||||
assert feature_store.keys() == [
|
||||
("node", None, "a"),
|
||||
("node", None, "b"),
|
||||
("node", None, "c"),
|
||||
]
|
||||
|
||||
|
||||
def test_basic_feature_store_hetero():
|
||||
a = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
b = torch.tensor([[[6], [8]], [[8], [9]]])
|
||||
metadata = {"max_value": 3}
|
||||
|
||||
features = {}
|
||||
features[("node", "author", "a")] = gb.TorchBasedFeature(
|
||||
a, metadata=metadata
|
||||
)
|
||||
features[("edge", "paper:cites", "b")] = gb.TorchBasedFeature(b)
|
||||
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
# Test __getitem__ to access the stored Feature.
|
||||
feature = feature_store[("node", "author", "a")]
|
||||
assert isinstance(feature, gb.Feature)
|
||||
assert torch.equal(
|
||||
feature.read(),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", "author", "a"),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_store.read("edge", "paper:cites", "b"),
|
||||
torch.tensor([[[6], [8]], [[8], [9]]]),
|
||||
)
|
||||
|
||||
# Test read with ids.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", "author", "a", torch.tensor([0])),
|
||||
torch.tensor([[1, 2, 4]]),
|
||||
)
|
||||
|
||||
# Test get the size of the entire feature.
|
||||
assert feature_store.size("node", "author", "a") == torch.Size([3])
|
||||
assert feature_store.size("edge", "paper:cites", "b") == torch.Size([2, 1])
|
||||
|
||||
# Test get metadata of the feature.
|
||||
assert feature_store.metadata("node", "author", "a") == metadata
|
||||
assert feature_store.metadata("edge", "paper:cites", "b") == {}
|
||||
|
||||
# Test __setitem__ and __contains__ of FeatureStore.
|
||||
assert ("node", "author", "c") not in feature_store
|
||||
feature_store[("node", "author", "c")] = feature_store[
|
||||
("node", "author", "a")
|
||||
]
|
||||
assert ("node", "author", "c") in feature_store
|
||||
|
||||
# Test get keys of the features.
|
||||
assert feature_store.keys() == [
|
||||
("node", "author", "a"),
|
||||
("edge", "paper:cites", "b"),
|
||||
("node", "author", "c"),
|
||||
]
|
||||
|
||||
|
||||
def test_basic_feature_store_errors():
|
||||
a = torch.tensor([3, 2, 1])
|
||||
b = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
|
||||
features = {}
|
||||
# Test error when dimension of the value is illegal.
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=rf"dimension of torch_feature in TorchBasedFeature must be "
|
||||
rf"greater than 1, but got {a.dim()} dimension.",
|
||||
):
|
||||
features[("node", "paper", "a")] = gb.TorchBasedFeature(a)
|
||||
features[("node", "author", "b")] = gb.TorchBasedFeature(b)
|
||||
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
# Test error when key does not exist.
|
||||
with pytest.raises(KeyError):
|
||||
feature_store.read("node", "paper", "b")
|
||||
|
||||
# Test error when at least one id is out of bound.
|
||||
with pytest.raises(IndexError):
|
||||
feature_store.read("node", "author", "b", torch.tensor([0, 3]))
|
||||
@@ -0,0 +1,68 @@
|
||||
import unittest
|
||||
|
||||
from functools import partial
|
||||
|
||||
import backend as F
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
WORLD_SIZE = 7
|
||||
|
||||
assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str != "gpu",
|
||||
reason="This test requires an NVIDIA GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("rank", list(range(WORLD_SIZE)))
|
||||
def test_rank_sort_and_unique_and_compact(dtype, rank):
|
||||
torch.manual_seed(7)
|
||||
nodes_list1 = [
|
||||
torch.randint(0, 2111111111, [777], dtype=dtype, device=F.ctx())
|
||||
for _ in range(10)
|
||||
]
|
||||
nodes_list2 = [nodes.sort()[0] for nodes in nodes_list1]
|
||||
|
||||
res1 = torch.ops.graphbolt.rank_sort(nodes_list1, rank, WORLD_SIZE)
|
||||
res2 = torch.ops.graphbolt.rank_sort(nodes_list2, rank, WORLD_SIZE)
|
||||
|
||||
for i, ((nodes1, idx1, offsets1), (nodes2, idx2, offsets2)) in enumerate(
|
||||
zip(res1, res2)
|
||||
):
|
||||
assert_equal(nodes_list1[i], nodes1[idx1])
|
||||
assert_equal(nodes_list2[i], nodes2[idx2])
|
||||
assert_equal(offsets1, offsets2)
|
||||
assert offsets1.is_pinned() and offsets2.is_pinned()
|
||||
|
||||
res3 = torch.ops.graphbolt.rank_sort(nodes_list1, rank, WORLD_SIZE)
|
||||
|
||||
# This function is deterministic. Call with identical arguments and check.
|
||||
for (nodes1, idx1, offsets1), (nodes3, idx3, offsets3) in zip(res1, res3):
|
||||
assert_equal(nodes1, nodes3)
|
||||
assert_equal(idx1, idx3)
|
||||
assert_equal(offsets1, offsets3)
|
||||
|
||||
# The dependency on the rank argument is simply a permutation.
|
||||
res4 = torch.ops.graphbolt.rank_sort(nodes_list1, 0, WORLD_SIZE)
|
||||
for (nodes1, idx1, offsets1), (nodes4, idx4, offsets4) in zip(res1, res4):
|
||||
off1 = offsets1.tolist()
|
||||
off4 = offsets4.tolist()
|
||||
assert_equal(nodes1[idx1], nodes4[idx4])
|
||||
for i in range(WORLD_SIZE):
|
||||
j = (i - rank + WORLD_SIZE) % WORLD_SIZE
|
||||
assert_equal(
|
||||
nodes1[off1[j] : off1[j + 1]], nodes4[off4[i] : off4[i + 1]]
|
||||
)
|
||||
|
||||
unique, compacted, offsets = gb.unique_and_compact(
|
||||
nodes_list1[:1], rank, WORLD_SIZE
|
||||
)
|
||||
|
||||
nodes1, idx1, offsets1 = res1[0]
|
||||
|
||||
assert_equal(unique, nodes1)
|
||||
assert_equal(compacted[0], idx1)
|
||||
assert_equal(offsets, offsets1)
|
||||
@@ -0,0 +1,184 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def to_on_disk_numpy(test_dir, name, t):
|
||||
path = os.path.join(test_dir, name + ".npy")
|
||||
np.save(path, t.numpy())
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("policy", ["s3-fifo", "sieve", "lru", "clock"])
|
||||
def test_cpu_cached_feature(dtype, policy):
|
||||
cache_size_a = 32
|
||||
cache_size_b = 64
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype)
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]], dtype=dtype)
|
||||
|
||||
pin_memory = F._default_context_str == "gpu"
|
||||
|
||||
cache_size_a *= a[:1].nbytes
|
||||
cache_size_b *= b[:1].nbytes
|
||||
|
||||
feat_store_a = gb.cpu_cached_feature(
|
||||
gb.TorchBasedFeature(a), cache_size_a, policy, pin_memory
|
||||
)
|
||||
feat_store_b = gb.cpu_cached_feature(
|
||||
gb.TorchBasedFeature(b), cache_size_b, policy, pin_memory
|
||||
)
|
||||
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(feat_store_a.read(), a)
|
||||
assert torch.equal(feat_store_b.read(), b)
|
||||
|
||||
# Test read with ids.
|
||||
assert torch.equal(
|
||||
# Test read when ids are on a different device.
|
||||
feat_store_a.read(torch.tensor([0], device=F.ctx())),
|
||||
torch.tensor([[1, 2, 3]], dtype=dtype, device=F.ctx()),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_b.read(torch.tensor([1, 1])),
|
||||
torch.tensor([[[4, 5], [6, 7]], [[4, 5], [6, 7]]], dtype=dtype),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_a.read(torch.tensor([1, 1])),
|
||||
torch.tensor([[4, 5, 6], [4, 5, 6]], dtype=dtype),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_b.read(torch.tensor([0])),
|
||||
torch.tensor([[[1, 2], [3, 4]]], dtype=dtype),
|
||||
)
|
||||
# The cache should be full now for the large cache sizes, %100 hit expected.
|
||||
total_miss = feat_store_a._feature.total_miss
|
||||
feat_store_a.read(torch.tensor([0, 1]))
|
||||
assert total_miss == feat_store_a._feature.total_miss
|
||||
total_miss = feat_store_b._feature.total_miss
|
||||
feat_store_b.read(torch.tensor([0, 1]))
|
||||
assert total_miss == feat_store_b._feature.total_miss
|
||||
assert feat_store_a._feature.miss_rate == feat_store_a.miss_rate
|
||||
|
||||
# Test get the size and count of the entire feature.
|
||||
assert feat_store_a.size() == torch.Size([3])
|
||||
assert feat_store_b.size() == torch.Size([2, 2])
|
||||
assert feat_store_a.count() == a.size(0)
|
||||
assert feat_store_b.count() == b.size(0)
|
||||
|
||||
# Test update the entire feature.
|
||||
feat_store_a.update(torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype))
|
||||
assert torch.equal(
|
||||
feat_store_a.read(),
|
||||
torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype),
|
||||
)
|
||||
|
||||
# Test update with ids.
|
||||
feat_store_a.update(
|
||||
torch.tensor([[2, 0, 1]], dtype=dtype),
|
||||
torch.tensor([0]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_a.read(),
|
||||
torch.tensor([[2, 0, 1], [3, 5, 2]], dtype=dtype),
|
||||
)
|
||||
|
||||
# Test with different dimensionality
|
||||
feat_store_a.update(b)
|
||||
assert torch.equal(feat_store_a.read(), b)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
def test_cpu_cached_feature_read_async(dtype):
|
||||
a = torch.randint(0, 2, [1000, 13], dtype=dtype)
|
||||
|
||||
cache_size = 256 * a[:1].nbytes
|
||||
|
||||
feat_store = gb.cpu_cached_feature(gb.TorchBasedFeature(a), cache_size)
|
||||
|
||||
# Test read with ids.
|
||||
ids1 = torch.tensor([0, 15, 71, 101])
|
||||
ids2 = torch.tensor([71, 101, 202, 303])
|
||||
for ids in [ids1, ids2]:
|
||||
reader = feat_store.read_async(ids)
|
||||
for _ in range(feat_store.read_async_num_stages(ids.device)):
|
||||
values = next(reader)
|
||||
assert torch.equal(values.wait(), a[ids])
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
not torch.ops.graphbolt.detect_io_uring(),
|
||||
reason="DiskBasedFeature is not available on this system.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
def test_cpu_cached_disk_feature_read_async(dtype):
|
||||
a = torch.randint(0, 2, [1000, 13], dtype=dtype)
|
||||
|
||||
cache_size = 256 * a[:1].nbytes
|
||||
|
||||
ids1 = torch.tensor([0, 15, 71, 101])
|
||||
ids2 = torch.tensor([71, 101, 202, 303])
|
||||
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
path = to_on_disk_numpy(test_dir, "tensor", a)
|
||||
|
||||
feat_store = gb.cpu_cached_feature(
|
||||
gb.DiskBasedFeature(path=path), cache_size
|
||||
)
|
||||
|
||||
# Test read feature.
|
||||
for ids in [ids1, ids2]:
|
||||
reader = feat_store.read_async(ids)
|
||||
for _ in range(feat_store.read_async_num_stages(ids.device)):
|
||||
values = next(reader)
|
||||
assert torch.equal(values.wait(), a[ids])
|
||||
|
||||
feat_store = None
|
||||
@@ -0,0 +1,192 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from functools import partial
|
||||
|
||||
import backend as F
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def to_on_disk_numpy(test_dir, name, t):
|
||||
path = os.path.join(test_dir, name + ".npy")
|
||||
t = t.numpy()
|
||||
np.save(path, t)
|
||||
return path
|
||||
|
||||
|
||||
assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
not torch.ops.graphbolt.detect_io_uring(),
|
||||
reason="DiskBasedFeature is not available on this system.",
|
||||
)
|
||||
def test_disk_based_feature():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
c = torch.randn([4111, 47])
|
||||
metadata = {"max_value": 3}
|
||||
path_a = to_on_disk_numpy(test_dir, "a", a)
|
||||
path_b = to_on_disk_numpy(test_dir, "b", b)
|
||||
path_c = to_on_disk_numpy(test_dir, "c", c)
|
||||
|
||||
feature_a = gb.DiskBasedFeature(path=path_a, metadata=metadata)
|
||||
feature_b = gb.DiskBasedFeature(path=path_b)
|
||||
feature_c = gb.DiskBasedFeature(path=path_c)
|
||||
|
||||
# Read the entire feature.
|
||||
assert_equal(feature_a.read(), torch.tensor([[1, 2, 3], [4, 5, 6]]))
|
||||
|
||||
assert_equal(
|
||||
feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
)
|
||||
|
||||
# Test read the feature with ids.
|
||||
assert_equal(
|
||||
feature_a.read(torch.tensor([0])),
|
||||
torch.tensor([[1, 2, 3]]),
|
||||
)
|
||||
assert_equal(
|
||||
feature_b.read(torch.tensor([1])),
|
||||
torch.tensor([[[4, 5], [6, 7]]]),
|
||||
)
|
||||
|
||||
# Test reading into pin_memory
|
||||
if F._default_context_str == "gpu":
|
||||
res = feature_a.read(torch.tensor([0], pin_memory=True))
|
||||
assert res.is_pinned()
|
||||
|
||||
# Test when the index tensor is large.
|
||||
torch_based_feature_a = gb.TorchBasedFeature(a)
|
||||
ind_a = torch.randint(low=0, high=a.size(0), size=(4111,))
|
||||
assert_equal(
|
||||
feature_a.read(ind_a),
|
||||
torch_based_feature_a.read(ind_a),
|
||||
)
|
||||
|
||||
# Test converting to torch_based_feature with read_into_memory()
|
||||
torch_based_feature_b = feature_b.read_into_memory()
|
||||
ind_b = torch.randint(low=0, high=b.size(0), size=(4111,))
|
||||
assert_equal(
|
||||
feature_b.read(ind_b),
|
||||
torch_based_feature_b.read(ind_b),
|
||||
)
|
||||
|
||||
# Test with larger stored feature tensor
|
||||
ind_c = torch.randint(low=0, high=c.size(0), size=(4111,))
|
||||
assert_equal(feature_c.read(ind_c), c[ind_c])
|
||||
|
||||
# Test get the size and count of the entire feature.
|
||||
assert feature_a.size() == torch.Size([3])
|
||||
assert feature_b.size() == torch.Size([2, 2])
|
||||
assert feature_a.count() == a.size(0)
|
||||
assert feature_b.count() == b.size(0)
|
||||
|
||||
# Test get metadata of the feature.
|
||||
assert feature_a.metadata() == metadata
|
||||
assert feature_b.metadata() == {}
|
||||
|
||||
with pytest.raises(IndexError):
|
||||
feature_a.read(torch.tensor([0, 1, 2, 3]))
|
||||
|
||||
# Test loading a Fortran contiguous ndarray.
|
||||
a_T = np.asfortranarray(a)
|
||||
path_a_T = test_dir + "a_T.npy"
|
||||
np.save(path_a_T, a_T)
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match="DiskBasedFeature only supports C_CONTIGUOUS array.",
|
||||
):
|
||||
gb.DiskBasedFeature(path=path_a_T, metadata=metadata)
|
||||
|
||||
# For windows, the file is locked by the numpy.load. We need to delete
|
||||
# it before closing the temporary directory.
|
||||
a = b = c = None
|
||||
feature_a = feature_b = feature_c = None
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
not torch.ops.graphbolt.detect_io_uring(),
|
||||
reason="DiskBasedFeature is not available on this system.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.int8,
|
||||
torch.float16,
|
||||
torch.complex128,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize(
|
||||
"shape", [(10, 20), (20, 10), (20, 25, 10), (137, 50, 30)]
|
||||
)
|
||||
@pytest.mark.parametrize("index", [[0], [1, 2, 3], [0, 6, 2, 8]])
|
||||
def test_more_disk_based_feature(dtype, idtype, shape, index):
|
||||
if dtype == torch.complex128:
|
||||
tensor = torch.complex(
|
||||
torch.randint(0, 127, shape, dtype=torch.float64),
|
||||
torch.randint(0, 127, shape, dtype=torch.float64),
|
||||
)
|
||||
else:
|
||||
tensor = torch.randint(0, 127, shape, dtype=dtype)
|
||||
test_tensor = tensor.clone()
|
||||
idx = torch.tensor(index, dtype=idtype)
|
||||
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
path = to_on_disk_numpy(test_dir, "tensor", tensor)
|
||||
|
||||
feature = gb.DiskBasedFeature(path=path)
|
||||
|
||||
# Test read feature.
|
||||
assert_equal(feature.read(idx), test_tensor[idx.long()])
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
not torch.ops.graphbolt.detect_io_uring(),
|
||||
reason="DiskBasedFeature is not available on this system.",
|
||||
)
|
||||
def test_disk_based_feature_repr():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
metadata = {"max_value": 3}
|
||||
|
||||
path_a = to_on_disk_numpy(test_dir, "a", a)
|
||||
path_b = to_on_disk_numpy(test_dir, "b", b)
|
||||
|
||||
feature_a = gb.DiskBasedFeature(path=path_a, metadata=metadata)
|
||||
feature_b = gb.DiskBasedFeature(path=path_b)
|
||||
|
||||
expected_str_feature_a = str(
|
||||
"DiskBasedFeature(\n"
|
||||
" feature=tensor([[1, 2, 3],\n"
|
||||
" [4, 5, 6]]),\n"
|
||||
" metadata={'max_value': 3},\n"
|
||||
")"
|
||||
)
|
||||
expected_str_feature_b = str(
|
||||
"DiskBasedFeature(\n"
|
||||
" feature=tensor([[[1, 2],\n"
|
||||
" [3, 4]],\n"
|
||||
"\n"
|
||||
" [[4, 5],\n"
|
||||
" [6, 7]]]),\n"
|
||||
" metadata={},\n"
|
||||
")"
|
||||
)
|
||||
assert str(feature_a) == expected_str_feature_a
|
||||
assert str(feature_b) == expected_str_feature_b
|
||||
a = b = metadata = None
|
||||
feature_a = feature_b = None
|
||||
expected_str_feature_a = expected_str_feature_b = None
|
||||
@@ -0,0 +1,172 @@
|
||||
import backend as F
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def _test_query_and_replace(policy1, policy2, keys, offset):
|
||||
# Testing query_and_replace equivalence to query and then replace.
|
||||
(
|
||||
_,
|
||||
index,
|
||||
pointers,
|
||||
missing_keys,
|
||||
found_offsets,
|
||||
missing_offsets,
|
||||
) = policy1.query_and_replace(keys, offset)
|
||||
found_cnt = keys.size(0) - missing_keys.size(0)
|
||||
found_pointers = pointers[:found_cnt]
|
||||
policy1.reading_completed(found_pointers, found_offsets)
|
||||
missing_pointers = pointers[found_cnt:]
|
||||
policy1.writing_completed(missing_pointers, missing_offsets)
|
||||
|
||||
(
|
||||
_,
|
||||
index2,
|
||||
missing_keys2,
|
||||
found_pointers2,
|
||||
found_offsets2,
|
||||
missing_offsets2,
|
||||
) = policy2.query(keys + offset, 0)
|
||||
policy2.reading_completed(found_pointers2, found_offsets2)
|
||||
(_, missing_pointers2, missing_offsets2) = policy2.replace(
|
||||
missing_keys2, missing_offsets2, 0
|
||||
)
|
||||
policy2.writing_completed(missing_pointers2, missing_offsets2)
|
||||
|
||||
assert torch.equal(index, index2)
|
||||
assert torch.equal(missing_keys, missing_keys2 - offset)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("offsets", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("feature_size", [2, 16])
|
||||
@pytest.mark.parametrize("num_parts", [1, 2, None])
|
||||
@pytest.mark.parametrize("policy", ["s3-fifo", "sieve", "lru", "clock"])
|
||||
@pytest.mark.parametrize("offset", [0, 1111111])
|
||||
def test_feature_cache(offsets, dtype, feature_size, num_parts, policy, offset):
|
||||
cache_size = 32 * (
|
||||
torch.get_num_threads() if num_parts is None else num_parts
|
||||
)
|
||||
a = torch.randint(0, 2, [1024, feature_size], dtype=dtype)
|
||||
cache = gb.impl.CPUFeatureCache(
|
||||
(cache_size,) + a.shape[1:], a.dtype, policy, num_parts
|
||||
)
|
||||
cache2 = gb.impl.CPUFeatureCache(
|
||||
(cache_size,) + a.shape[1:], a.dtype, policy, num_parts
|
||||
)
|
||||
policy1 = gb.impl.CPUFeatureCache(
|
||||
(cache_size,) + a.shape[1:], a.dtype, policy, num_parts
|
||||
)._policy
|
||||
policy2 = gb.impl.CPUFeatureCache(
|
||||
(cache_size,) + a.shape[1:], a.dtype, policy, num_parts
|
||||
)._policy
|
||||
reader_fn = lambda keys: a[keys]
|
||||
|
||||
keys = torch.tensor([0, 1])
|
||||
values, missing_index, missing_keys, missing_offsets = cache.query(
|
||||
keys, offset
|
||||
)
|
||||
if not offsets:
|
||||
missing_offsets = None
|
||||
assert torch.equal(
|
||||
missing_keys.flip([0]) if num_parts == 1 else missing_keys.sort()[0],
|
||||
keys,
|
||||
)
|
||||
|
||||
missing_values = a[missing_keys]
|
||||
cache.replace(missing_keys, missing_values, missing_offsets, offset)
|
||||
values[missing_index] = missing_values
|
||||
assert torch.equal(values, a[keys])
|
||||
assert torch.equal(
|
||||
cache2.query_and_replace(keys, reader_fn, offset), a[keys]
|
||||
)
|
||||
|
||||
_test_query_and_replace(policy1, policy2, keys, offset)
|
||||
|
||||
pin_memory = F._default_context_str == "gpu"
|
||||
|
||||
keys = torch.arange(1, 33, pin_memory=pin_memory)
|
||||
values, missing_index, missing_keys, missing_offsets = cache.query(
|
||||
keys, offset
|
||||
)
|
||||
if not offsets:
|
||||
missing_offsets = None
|
||||
assert torch.equal(
|
||||
missing_keys.flip([0]) if num_parts == 1 else missing_keys.sort()[0],
|
||||
torch.arange(2, 33),
|
||||
)
|
||||
assert not pin_memory or values.is_pinned()
|
||||
|
||||
missing_values = a[missing_keys]
|
||||
cache.replace(missing_keys, missing_values, missing_offsets, offset)
|
||||
values[missing_index] = missing_values
|
||||
assert torch.equal(values, a[keys])
|
||||
assert torch.equal(
|
||||
cache2.query_and_replace(keys, reader_fn, offset), a[keys]
|
||||
)
|
||||
|
||||
_test_query_and_replace(policy1, policy2, keys, offset)
|
||||
|
||||
values, missing_index, missing_keys, missing_offsets = cache.query(
|
||||
keys, offset
|
||||
)
|
||||
if not offsets:
|
||||
missing_offsets = None
|
||||
assert torch.equal(missing_keys.flip([0]), torch.tensor([]))
|
||||
|
||||
missing_values = a[missing_keys]
|
||||
cache.replace(missing_keys, missing_values, missing_offsets, offset)
|
||||
values[missing_index] = missing_values
|
||||
assert torch.equal(values, a[keys])
|
||||
assert torch.equal(
|
||||
cache2.query_and_replace(keys, reader_fn, offset), a[keys]
|
||||
)
|
||||
|
||||
_test_query_and_replace(policy1, policy2, keys, offset)
|
||||
|
||||
values, missing_index, missing_keys, missing_offsets = cache.query(
|
||||
keys, offset
|
||||
)
|
||||
if not offsets:
|
||||
missing_offsets = None
|
||||
assert torch.equal(missing_keys.flip([0]), torch.tensor([]))
|
||||
|
||||
missing_values = a[missing_keys]
|
||||
cache.replace(missing_keys, missing_values, missing_offsets, offset)
|
||||
values[missing_index] = missing_values
|
||||
assert torch.equal(values, a[keys])
|
||||
assert torch.equal(
|
||||
cache2.query_and_replace(keys, reader_fn, offset), a[keys]
|
||||
)
|
||||
|
||||
_test_query_and_replace(policy1, policy2, keys, offset)
|
||||
|
||||
assert cache.miss_rate == cache2.miss_rate
|
||||
|
||||
raw_feature_cache = torch.ops.graphbolt.feature_cache(
|
||||
(cache_size,) + a.shape[1:], a.dtype, pin_memory
|
||||
)
|
||||
idx = torch.tensor([0, 1, 2])
|
||||
raw_feature_cache.replace(idx, a[idx])
|
||||
val = raw_feature_cache.index_select(idx)
|
||||
assert torch.equal(val, a[idx])
|
||||
if pin_memory:
|
||||
val = raw_feature_cache.index_select(idx.to(F.ctx()))
|
||||
assert torch.equal(val, a[idx].to(F.ctx()))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,219 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def to_on_disk_numpy(test_dir, name, t):
|
||||
path = os.path.join(test_dir, name + ".npy")
|
||||
np.save(path, t.cpu().numpy())
|
||||
return path
|
||||
|
||||
|
||||
def _skip_condition_cached_feature():
|
||||
return (F._default_context_str != "gpu") or (
|
||||
torch.cuda.get_device_capability()[0] < 7
|
||||
)
|
||||
|
||||
|
||||
def _reason_to_skip_cached_feature():
|
||||
if F._default_context_str != "gpu":
|
||||
return "GPUCachedFeature tests are available only when testing the GPU backend."
|
||||
|
||||
return "GPUCachedFeature requires a Volta or later generation NVIDIA GPU."
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
_skip_condition_cached_feature(),
|
||||
reason=_reason_to_skip_cached_feature(),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("cache_size_a", [1, 1024])
|
||||
@pytest.mark.parametrize("cache_size_b", [1, 1024])
|
||||
def test_gpu_cached_feature(dtype, cache_size_a, cache_size_b):
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=dtype, pin_memory=True)
|
||||
b = torch.tensor(
|
||||
[[[1, 2], [3, 4]], [[4, 5], [6, 7]]], dtype=dtype, pin_memory=True
|
||||
)
|
||||
|
||||
cache_size_a *= a[:1].element_size() * a[:1].numel()
|
||||
cache_size_b *= b[:1].element_size() * b[:1].numel()
|
||||
|
||||
feat_store_a = gb.gpu_cached_feature(gb.TorchBasedFeature(a), cache_size_a)
|
||||
feat_store_b = gb.gpu_cached_feature(gb.TorchBasedFeature(b), cache_size_b)
|
||||
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(feat_store_a.read(), a.to("cuda"))
|
||||
assert torch.equal(feat_store_b.read(), b.to("cuda"))
|
||||
|
||||
# Test read with ids.
|
||||
assert torch.equal(
|
||||
feat_store_a.read(torch.tensor([0]).to("cuda")),
|
||||
torch.tensor([[1, 2, 3]], dtype=dtype).to("cuda"),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_b.read(torch.tensor([1, 1]).to("cuda")),
|
||||
torch.tensor([[[4, 5], [6, 7]], [[4, 5], [6, 7]]], dtype=dtype).to(
|
||||
"cuda"
|
||||
),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_a.read(torch.tensor([1, 1]).to("cuda")),
|
||||
torch.tensor([[4, 5, 6], [4, 5, 6]], dtype=dtype).to("cuda"),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_b.read(torch.tensor([0]).to("cuda")),
|
||||
torch.tensor([[[1, 2], [3, 4]]], dtype=dtype).to("cuda"),
|
||||
)
|
||||
# The cache should be full now for the large cache sizes, %100 hit expected.
|
||||
if cache_size_a >= 1024:
|
||||
total_miss = feat_store_a._feature.total_miss
|
||||
feat_store_a.read(torch.tensor([0, 1]).to("cuda"))
|
||||
assert total_miss == feat_store_a._feature.total_miss
|
||||
if cache_size_b >= 1024:
|
||||
total_miss = feat_store_b._feature.total_miss
|
||||
feat_store_b.read(torch.tensor([0, 1]).to("cuda"))
|
||||
assert total_miss == feat_store_b._feature.total_miss
|
||||
assert feat_store_a._feature.miss_rate == feat_store_a.miss_rate
|
||||
|
||||
# Test get the size and count of the entire feature.
|
||||
assert feat_store_a.size() == torch.Size([3])
|
||||
assert feat_store_b.size() == torch.Size([2, 2])
|
||||
assert feat_store_a.count() == a.size(0)
|
||||
assert feat_store_b.count() == b.size(0)
|
||||
|
||||
# Test update the entire feature.
|
||||
feat_store_a.update(
|
||||
torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype).to("cuda")
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_a.read(),
|
||||
torch.tensor([[0, 1, 2], [3, 5, 2]], dtype=dtype).to("cuda"),
|
||||
)
|
||||
|
||||
# Test update with ids.
|
||||
feat_store_a.update(
|
||||
torch.tensor([[2, 0, 1]], dtype=dtype).to("cuda"),
|
||||
torch.tensor([0]).to("cuda"),
|
||||
)
|
||||
assert torch.equal(
|
||||
feat_store_a.read(),
|
||||
torch.tensor([[2, 0, 1], [3, 5, 2]], dtype=dtype).to("cuda"),
|
||||
)
|
||||
|
||||
# Test with different dimensionality
|
||||
feat_store_a.update(b)
|
||||
assert torch.equal(feat_store_a.read(), b.to("cuda"))
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
_skip_condition_cached_feature(),
|
||||
reason=_reason_to_skip_cached_feature(),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("pin_memory", [False, True])
|
||||
def test_gpu_cached_feature_read_async(dtype, pin_memory):
|
||||
a = torch.randint(0, 2, [1000, 13], dtype=dtype, pin_memory=pin_memory)
|
||||
a_cuda = a.to(F.ctx())
|
||||
|
||||
cache_size = 256 * a[:1].nbytes
|
||||
|
||||
feat_store = gb.gpu_cached_feature(gb.TorchBasedFeature(a), cache_size)
|
||||
|
||||
# Test read with ids.
|
||||
ids1 = torch.tensor([0, 15, 71, 101], device=F.ctx())
|
||||
ids2 = torch.tensor([71, 101, 202, 303], device=F.ctx())
|
||||
for ids in [ids1, ids2]:
|
||||
reader = feat_store.read_async(ids)
|
||||
for _ in range(feat_store.read_async_num_stages(ids.device)):
|
||||
values = next(reader)
|
||||
assert torch.equal(values.wait(), a_cuda[ids])
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
_skip_condition_cached_feature(),
|
||||
reason=_reason_to_skip_cached_feature(),
|
||||
)
|
||||
@unittest.skipIf(
|
||||
not torch.ops.graphbolt.detect_io_uring(),
|
||||
reason="DiskBasedFeature is not available on this system.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
def test_gpu_cached_nested_feature_async(dtype):
|
||||
a = torch.randint(0, 2, [1000, 13], dtype=dtype, device=F.ctx())
|
||||
|
||||
cache_size = 256 * a[:1].nbytes
|
||||
|
||||
ids1 = torch.tensor([0, 15, 71, 101], device=F.ctx())
|
||||
ids2 = torch.tensor([71, 101, 202, 303], device=F.ctx())
|
||||
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
path = to_on_disk_numpy(test_dir, "tensor", a)
|
||||
|
||||
disk_store = gb.DiskBasedFeature(path=path)
|
||||
feat_store1 = gb.gpu_cached_feature(disk_store, cache_size)
|
||||
feat_store2 = gb.gpu_cached_feature(
|
||||
gb.cpu_cached_feature(disk_store, cache_size * 2), cache_size
|
||||
)
|
||||
feat_store3 = gb.gpu_cached_feature(
|
||||
gb.cpu_cached_feature(disk_store, cache_size * 2, pin_memory=True),
|
||||
cache_size,
|
||||
)
|
||||
|
||||
# Test read feature.
|
||||
for feat_store in [feat_store1, feat_store2, feat_store3]:
|
||||
for ids in [ids1, ids2]:
|
||||
reader = feat_store.read_async(ids)
|
||||
for _ in range(feat_store.read_async_num_stages(ids.device)):
|
||||
values = next(reader)
|
||||
assert torch.equal(values.wait(), a[ids])
|
||||
|
||||
feat_store1 = feat_store2 = feat_store3 = disk_store = None
|
||||
@@ -0,0 +1,82 @@
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str != "gpu"
|
||||
or torch.cuda.get_device_capability()[0] < 7,
|
||||
reason="GPUCachedFeature tests are available only when testing the GPU backend."
|
||||
if F._default_context_str != "gpu"
|
||||
else "GPUCachedFeature requires a Volta or later generation NVIDIA GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"indptr_dtype",
|
||||
[
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("cache_size", [4, 9, 11])
|
||||
@pytest.mark.parametrize("with_edge_ids", [True, False])
|
||||
def test_gpu_graph_cache(indptr_dtype, dtype, cache_size, with_edge_ids):
|
||||
indices_dtype = torch.int32
|
||||
indptr = torch.tensor([0, 3, 6, 10], dtype=indptr_dtype, pin_memory=True)
|
||||
indices = torch.arange(0, indptr[-1], dtype=indices_dtype, pin_memory=True)
|
||||
probs_or_mask = indices.to(dtype).pin_memory()
|
||||
edge_tensors = [indices, probs_or_mask]
|
||||
|
||||
g = gb.GPUGraphCache(
|
||||
cache_size,
|
||||
2,
|
||||
indptr.dtype,
|
||||
[e.dtype for e in edge_tensors],
|
||||
not with_edge_ids,
|
||||
)
|
||||
|
||||
for i in range(10):
|
||||
keys = (
|
||||
torch.arange(2, dtype=indices_dtype, device=F.ctx()) + i * 2
|
||||
) % (indptr.size(0) - 1)
|
||||
missing_keys, replace = g.query(keys)
|
||||
(
|
||||
missing_indptr,
|
||||
missing_edge_tensors,
|
||||
) = torch.ops.graphbolt.index_select_csc_batched(
|
||||
indptr, edge_tensors, missing_keys, with_edge_ids, None
|
||||
)
|
||||
output_indptr, output_edge_tensors = replace(
|
||||
missing_indptr, missing_edge_tensors
|
||||
)
|
||||
|
||||
(
|
||||
reference_indptr,
|
||||
reference_edge_tensors,
|
||||
) = torch.ops.graphbolt.index_select_csc_batched(
|
||||
indptr, edge_tensors, keys, with_edge_ids, None
|
||||
)
|
||||
|
||||
assert torch.equal(output_indptr, reference_indptr)
|
||||
assert len(output_edge_tensors) == len(reference_edge_tensors)
|
||||
for e, ref in zip(output_edge_tensors, reference_edge_tensors):
|
||||
assert torch.equal(e, ref)
|
||||
@@ -0,0 +1,42 @@
|
||||
import backend as F
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cached_feature_type", [gb.cpu_cached_feature, gb.gpu_cached_feature]
|
||||
)
|
||||
def test_hetero_cached_feature(cached_feature_type):
|
||||
if cached_feature_type == gb.gpu_cached_feature and (
|
||||
F._default_context_str != "gpu"
|
||||
or torch.cuda.get_device_capability()[0] < 7
|
||||
):
|
||||
pytest.skip(
|
||||
"GPUCachedFeature tests are available only when testing the GPU backend."
|
||||
if F._default_context_str != "gpu"
|
||||
else "GPUCachedFeature requires a Volta or later generation NVIDIA GPU."
|
||||
)
|
||||
device = F.ctx() if cached_feature_type == gb.gpu_cached_feature else None
|
||||
pin_memory = cached_feature_type == gb.gpu_cached_feature
|
||||
|
||||
a = {
|
||||
("node", str(i), "feat"): gb.TorchBasedFeature(
|
||||
torch.randn([(i + 1) * 10, 5], pin_memory=pin_memory)
|
||||
)
|
||||
for i in range(75)
|
||||
}
|
||||
cached_a = cached_feature_type(a, 2**18)
|
||||
|
||||
for i in range(1024):
|
||||
etype = i % len(a)
|
||||
ids = torch.randint(
|
||||
0, (etype + 1) * 10 - 1, ((etype + 1) * 4,), device=device
|
||||
)
|
||||
feature_key = ("node", str(etype), "feat")
|
||||
ref = a[feature_key].read(ids)
|
||||
val = cached_a[feature_key].read(ids)
|
||||
torch.testing.assert_close(ref, val, rtol=0, atol=0)
|
||||
assert cached_a[feature_key].miss_rate < 0.69
|
||||
@@ -0,0 +1,287 @@
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from .. import gb_test_utils
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "cpu",
|
||||
reason="Tests for pinned memory are only meaningful on GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"indptr_dtype",
|
||||
[torch.int32, torch.int64],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"indices_dtype",
|
||||
[
|
||||
torch.int8,
|
||||
torch.uint8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("is_pinned", [False, True])
|
||||
@pytest.mark.parametrize("with_edge_ids", [False, True])
|
||||
@pytest.mark.parametrize("output_size", [None, True])
|
||||
def test_index_select_csc(
|
||||
indptr_dtype, indices_dtype, idtype, is_pinned, with_edge_ids, output_size
|
||||
):
|
||||
"""Original graph in COO:
|
||||
1 0 1 0 1 0
|
||||
1 0 0 1 0 1
|
||||
0 1 0 1 0 0
|
||||
0 1 0 0 1 0
|
||||
1 0 0 0 0 1
|
||||
0 0 1 0 1 0
|
||||
"""
|
||||
indptr = torch.tensor([0, 3, 5, 7, 9, 12, 14], dtype=indptr_dtype)
|
||||
indices = torch.tensor(
|
||||
[0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4], dtype=indices_dtype
|
||||
)
|
||||
index = torch.tensor([0, 5, 3], dtype=idtype)
|
||||
|
||||
cpu_indptr, cpu_indices = torch.ops.graphbolt.index_select_csc(
|
||||
indptr, indices, index, None
|
||||
)
|
||||
if is_pinned:
|
||||
indptr = indptr.pin_memory()
|
||||
indices = indices.pin_memory()
|
||||
else:
|
||||
indptr = indptr.cuda()
|
||||
indices = indices.cuda()
|
||||
index = index.cuda()
|
||||
edge_ids = torch.tensor(
|
||||
[0, 1, 2, 12, 13, 7, 8], dtype=indptr_dtype, device=index.device
|
||||
)
|
||||
|
||||
if output_size:
|
||||
output_size = len(cpu_indices)
|
||||
|
||||
gpu_indptr, gpu_indices = torch.ops.graphbolt.index_select_csc(
|
||||
indptr, indices, index, output_size
|
||||
)
|
||||
assert not cpu_indptr.is_cuda
|
||||
assert not cpu_indices.is_cuda
|
||||
|
||||
assert gpu_indptr.is_cuda
|
||||
assert gpu_indices.is_cuda
|
||||
|
||||
assert torch.equal(cpu_indptr, gpu_indptr.cpu())
|
||||
assert torch.equal(cpu_indices, gpu_indices.cpu())
|
||||
|
||||
for output_size_selection in [None, output_size]:
|
||||
indices_list = [
|
||||
indices,
|
||||
indices.int().pin_memory() if is_pinned else indices.int(),
|
||||
]
|
||||
(
|
||||
gpu_indptr2,
|
||||
gpu_indices_list,
|
||||
) = torch.ops.graphbolt.index_select_csc_batched(
|
||||
indptr, indices_list, index, with_edge_ids, output_size_selection
|
||||
)
|
||||
|
||||
assert torch.equal(gpu_indptr, gpu_indptr2)
|
||||
assert torch.equal(gpu_indices_list[0], gpu_indices)
|
||||
assert torch.equal(gpu_indices_list[1], gpu_indices.int())
|
||||
if with_edge_ids:
|
||||
assert torch.equal(gpu_indices_list[2], edge_ids)
|
||||
|
||||
|
||||
def test_InSubgraphSampler_homo():
|
||||
"""Original graph in COO:
|
||||
1 0 1 0 1 0
|
||||
1 0 0 1 0 1
|
||||
0 1 0 1 0 0
|
||||
0 1 0 0 1 0
|
||||
1 0 0 0 0 1
|
||||
0 0 1 0 1 0
|
||||
"""
|
||||
indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
|
||||
indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
|
||||
graph = gb.fused_csc_sampling_graph(indptr, indices).to(F.ctx())
|
||||
|
||||
seed_nodes = torch.LongTensor([0, 5, 3])
|
||||
item_set = gb.ItemSet(seed_nodes, names="seeds")
|
||||
batch_size = 1
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
|
||||
in_subgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
|
||||
|
||||
it = iter(in_subgraph_sampler)
|
||||
|
||||
def original_indices(minibatch):
|
||||
sampled_subgraph = minibatch.sampled_subgraphs[0]
|
||||
_indices = sampled_subgraph.original_row_node_ids[
|
||||
sampled_subgraph.sampled_csc.indices
|
||||
]
|
||||
return _indices
|
||||
|
||||
mn = next(it)
|
||||
assert torch.equal(mn.seeds, torch.LongTensor([0]).to(F.ctx()))
|
||||
assert torch.equal(
|
||||
mn.sampled_subgraphs[0].sampled_csc.indptr,
|
||||
torch.tensor([0, 3]).to(F.ctx()),
|
||||
)
|
||||
|
||||
mn = next(it)
|
||||
assert torch.equal(mn.seeds, torch.LongTensor([5]).to(F.ctx()))
|
||||
assert torch.equal(
|
||||
mn.sampled_subgraphs[0].sampled_csc.indptr,
|
||||
torch.tensor([0, 2]).to(F.ctx()),
|
||||
)
|
||||
assert torch.equal(original_indices(mn), torch.tensor([1, 4]).to(F.ctx()))
|
||||
|
||||
mn = next(it)
|
||||
assert torch.equal(mn.seeds, torch.LongTensor([3]).to(F.ctx()))
|
||||
assert torch.equal(
|
||||
mn.sampled_subgraphs[0].sampled_csc.indptr,
|
||||
torch.tensor([0, 2]).to(F.ctx()),
|
||||
)
|
||||
assert torch.equal(original_indices(mn), torch.tensor([1, 2]).to(F.ctx()))
|
||||
|
||||
|
||||
def test_InSubgraphSampler_hetero():
|
||||
"""Original graph in COO:
|
||||
1 0 1 0 1 0
|
||||
1 0 0 1 0 1
|
||||
0 1 0 1 0 0
|
||||
0 1 0 0 1 0
|
||||
1 0 0 0 0 1
|
||||
0 0 1 0 1 0
|
||||
node_type_0: [0, 1, 2]
|
||||
node_type_1: [3, 4, 5]
|
||||
edge_type_0: node_type_0 -> node_type_0
|
||||
edge_type_1: node_type_0 -> node_type_1
|
||||
edge_type_2: node_type_1 -> node_type_0
|
||||
edge_type_3: node_type_1 -> node_type_1
|
||||
"""
|
||||
ntypes = {
|
||||
"N0": 0,
|
||||
"N1": 1,
|
||||
}
|
||||
etypes = {
|
||||
"N0:R0:N0": 0,
|
||||
"N0:R1:N1": 1,
|
||||
"N1:R2:N0": 2,
|
||||
"N1:R3:N1": 3,
|
||||
}
|
||||
indptr = torch.LongTensor([0, 3, 5, 7, 9, 12, 14])
|
||||
indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 5, 1, 2, 0, 3, 5, 1, 4])
|
||||
node_type_offset = torch.LongTensor([0, 3, 6])
|
||||
type_per_edge = torch.LongTensor([0, 0, 2, 0, 2, 0, 2, 1, 1, 1, 3, 3, 1, 3])
|
||||
graph = gb.fused_csc_sampling_graph(
|
||||
csc_indptr=indptr,
|
||||
indices=indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
).to(F.ctx())
|
||||
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"N0": gb.ItemSet(torch.LongTensor([1, 0, 2]), names="seeds"),
|
||||
"N1": gb.ItemSet(torch.LongTensor([0, 2, 1]), names="seeds"),
|
||||
}
|
||||
)
|
||||
batch_size = 2
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
|
||||
in_subgraph_sampler = gb.InSubgraphSampler(item_sampler, graph)
|
||||
|
||||
it = iter(in_subgraph_sampler)
|
||||
|
||||
mn = next(it)
|
||||
assert torch.equal(mn.seeds["N0"], torch.LongTensor([1, 0]).to(F.ctx()))
|
||||
expected_sampled_csc = {
|
||||
"N0:R0:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1, 3]),
|
||||
indices=torch.LongTensor([2, 1, 0]),
|
||||
),
|
||||
"N0:R1:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
|
||||
),
|
||||
"N1:R2:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([0, 1])
|
||||
),
|
||||
"N1:R3:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
|
||||
),
|
||||
}
|
||||
for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
|
||||
assert torch.equal(
|
||||
pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
|
||||
)
|
||||
assert torch.equal(
|
||||
pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
|
||||
)
|
||||
|
||||
mn = next(it)
|
||||
assert mn.seeds == {
|
||||
"N0": torch.LongTensor([2]).to(F.ctx()),
|
||||
"N1": torch.LongTensor([0]).to(F.ctx()),
|
||||
}
|
||||
expected_sampled_csc = {
|
||||
"N0:R0:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1]), indices=torch.LongTensor([1])
|
||||
),
|
||||
"N0:R1:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 2]), indices=torch.LongTensor([2, 0])
|
||||
),
|
||||
"N1:R2:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1]), indices=torch.LongTensor([1])
|
||||
),
|
||||
"N1:R3:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 0]), indices=torch.LongTensor([])
|
||||
),
|
||||
}
|
||||
for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
|
||||
assert torch.equal(
|
||||
pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
|
||||
)
|
||||
assert torch.equal(
|
||||
pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
|
||||
)
|
||||
|
||||
mn = next(it)
|
||||
assert torch.equal(mn.seeds["N1"], torch.LongTensor([2, 1]).to(F.ctx()))
|
||||
expected_sampled_csc = {
|
||||
"N0:R0:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
|
||||
),
|
||||
"N0:R1:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([0, 1])
|
||||
),
|
||||
"N1:R2:N0": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0]), indices=torch.LongTensor([])
|
||||
),
|
||||
"N1:R3:N1": gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1, 3]),
|
||||
indices=torch.LongTensor([1, 2, 0]),
|
||||
),
|
||||
}
|
||||
if graph.csc_indptr.is_cuda and torch.cuda.get_device_capability()[0] < 7:
|
||||
expected_sampled_csc["N0:R1:N1"] = gb.CSCFormatBase(
|
||||
indptr=torch.LongTensor([0, 1, 2]), indices=torch.LongTensor([1, 0])
|
||||
)
|
||||
for etype, pairs in mn.sampled_subgraphs[0].sampled_csc.items():
|
||||
assert torch.equal(
|
||||
pairs.indices, expected_sampled_csc[etype].indices.to(F.ctx())
|
||||
)
|
||||
assert torch.equal(
|
||||
pairs.indptr, expected_sampled_csc[etype].indptr.to(F.ctx())
|
||||
)
|
||||
@@ -0,0 +1,36 @@
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
from dgl import AddSelfLoop
|
||||
from dgl.data import AsNodePredDataset, CoraGraphDataset
|
||||
|
||||
|
||||
def test_LegacyDataset_homo_node_pred():
|
||||
cora = CoraGraphDataset(transform=AddSelfLoop())
|
||||
dataset = gb.LegacyDataset(cora)
|
||||
|
||||
# Check tasks.
|
||||
assert len(dataset.tasks) == 1
|
||||
task = dataset.tasks[0]
|
||||
assert task.train_set.names == ("seeds", "labels")
|
||||
assert len(task.train_set) == 140
|
||||
assert task.validation_set.names == ("seeds", "labels")
|
||||
assert len(task.validation_set) == 500
|
||||
assert task.test_set.names == ("seeds", "labels")
|
||||
assert len(task.test_set) == 1000
|
||||
assert task.metadata["num_classes"] == 7
|
||||
|
||||
num_nodes = 2708
|
||||
assert dataset.graph.num_nodes == num_nodes
|
||||
assert len(dataset.all_nodes_set) == num_nodes
|
||||
assert dataset.feature.size("node", None, "feat") == torch.Size([1433])
|
||||
assert (
|
||||
dataset.feature.read(
|
||||
"node", None, "feat", torch.tensor([num_nodes - 1])
|
||||
).size(dim=0)
|
||||
== 1
|
||||
)
|
||||
# Out of bound indexing results in segmentation fault instead of exception
|
||||
# in CI. This may be related to docker env. Skip it for now.
|
||||
# with pytest.raises(IndexError):
|
||||
# dataset.feature.read("node", None, "feat", torch.Tensor([num_nodes]))
|
||||
@@ -0,0 +1,257 @@
|
||||
import re
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from .. import gb_test_utils
|
||||
|
||||
|
||||
def test_NegativeSampler_invoke():
|
||||
# Instantiate graph and required datapipes.
|
||||
num_seeds = 30
|
||||
item_set = gb.ItemSet(
|
||||
torch.arange(0, 2 * num_seeds).reshape(-1, 2), names="seeds"
|
||||
)
|
||||
batch_size = 10
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
negative_ratio = 2
|
||||
|
||||
# Invoke NegativeSampler via class constructor.
|
||||
negative_sampler = gb.NegativeSampler(
|
||||
item_sampler,
|
||||
negative_ratio,
|
||||
)
|
||||
with pytest.raises(NotImplementedError):
|
||||
next(iter(negative_sampler))
|
||||
|
||||
# Invoke NegativeSampler via functional form.
|
||||
negative_sampler = item_sampler.sample_negative(
|
||||
negative_ratio,
|
||||
)
|
||||
with pytest.raises(NotImplementedError):
|
||||
next(iter(negative_sampler))
|
||||
|
||||
|
||||
def test_UniformNegativeSampler_invoke():
|
||||
# Instantiate graph and required datapipes.
|
||||
graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
|
||||
F.ctx()
|
||||
)
|
||||
num_seeds = 30
|
||||
item_set = gb.ItemSet(
|
||||
torch.arange(0, 2 * num_seeds).reshape(-1, 2), names="seeds"
|
||||
)
|
||||
batch_size = 10
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
negative_ratio = 2
|
||||
|
||||
def _verify(negative_sampler):
|
||||
for data in negative_sampler:
|
||||
# Assertation
|
||||
seeds_len = batch_size + batch_size * negative_ratio
|
||||
assert data.seeds.size(0) == seeds_len
|
||||
assert data.labels.size(0) == seeds_len
|
||||
assert data.indexes.size(0) == seeds_len
|
||||
|
||||
# Invoke UniformNegativeSampler via class constructor.
|
||||
negative_sampler = gb.UniformNegativeSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
_verify(negative_sampler)
|
||||
|
||||
# Invoke UniformNegativeSampler via functional form.
|
||||
negative_sampler = item_sampler.sample_uniform_negative(
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
_verify(negative_sampler)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("negative_ratio", [1, 5, 10, 20])
|
||||
def test_Uniform_NegativeSampler(negative_ratio):
|
||||
# Construct FusedCSCSamplingGraph.
|
||||
graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
|
||||
F.ctx()
|
||||
)
|
||||
num_seeds = 30
|
||||
item_set = gb.ItemSet(
|
||||
torch.arange(0, num_seeds * 2).reshape(-1, 2), names="seeds"
|
||||
)
|
||||
batch_size = 10
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
# Construct NegativeSampler.
|
||||
negative_sampler = gb.UniformNegativeSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
# Perform Negative sampling.
|
||||
for data in negative_sampler:
|
||||
seeds_len = batch_size + batch_size * negative_ratio
|
||||
# Assertation
|
||||
assert data.seeds.size(0) == seeds_len
|
||||
assert data.labels.size(0) == seeds_len
|
||||
assert data.indexes.size(0) == seeds_len
|
||||
# Check negative seeds value.
|
||||
pos_src = data.seeds[:batch_size, 0]
|
||||
neg_src = data.seeds[batch_size:, 0]
|
||||
assert torch.equal(pos_src.repeat_interleave(negative_ratio), neg_src)
|
||||
# Check labels.
|
||||
assert torch.equal(
|
||||
data.labels[:batch_size], torch.ones(batch_size).to(F.ctx())
|
||||
)
|
||||
assert torch.equal(
|
||||
data.labels[batch_size:],
|
||||
torch.zeros(batch_size * negative_ratio).to(F.ctx()),
|
||||
)
|
||||
# Check indexes.
|
||||
pos_indexes = torch.arange(0, batch_size).to(F.ctx())
|
||||
neg_indexes = pos_indexes.repeat_interleave(negative_ratio)
|
||||
expected_indexes = torch.cat((pos_indexes, neg_indexes))
|
||||
assert torch.equal(data.indexes, expected_indexes)
|
||||
|
||||
|
||||
def test_Uniform_NegativeSampler_error_shape():
|
||||
# 1. seeds with shape N*3.
|
||||
# Construct FusedCSCSamplingGraph.
|
||||
graph = gb_test_utils.rand_csc_graph(100, 0.05, bidirection_edge=True).to(
|
||||
F.ctx()
|
||||
)
|
||||
num_seeds = 30
|
||||
item_set = gb.ItemSet(
|
||||
torch.arange(0, num_seeds * 3).reshape(-1, 3), names="seeds"
|
||||
)
|
||||
batch_size = 10
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
negative_ratio = 2
|
||||
# Construct NegativeSampler.
|
||||
negative_sampler = gb.UniformNegativeSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Only tensor with shape N*2 is "
|
||||
+ "supported for negative sampling, but got torch.Size([10, 3])."
|
||||
),
|
||||
):
|
||||
next(iter(negative_sampler))
|
||||
|
||||
# 2. seeds with shape N*2*1.
|
||||
# Construct FusedCSCSamplingGraph.
|
||||
item_set = gb.ItemSet(
|
||||
torch.arange(0, num_seeds * 2).reshape(-1, 2, 1), names="seeds"
|
||||
)
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
# Construct NegativeSampler.
|
||||
negative_sampler = gb.UniformNegativeSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Only tensor with shape N*2 is "
|
||||
+ "supported for negative sampling, but got torch.Size([10, 2, 1])."
|
||||
),
|
||||
):
|
||||
next(iter(negative_sampler))
|
||||
|
||||
# 3. seeds with shape N.
|
||||
# Construct FusedCSCSamplingGraph.
|
||||
item_set = gb.ItemSet(torch.arange(0, num_seeds), names="seeds")
|
||||
item_sampler = gb.ItemSampler(item_set, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
# Construct NegativeSampler.
|
||||
negative_sampler = gb.UniformNegativeSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
negative_ratio,
|
||||
)
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Only tensor with shape N*2 is "
|
||||
+ "supported for negative sampling, but got torch.Size([10])."
|
||||
),
|
||||
):
|
||||
next(iter(negative_sampler))
|
||||
|
||||
|
||||
def get_hetero_graph():
|
||||
# COO graph:
|
||||
# [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
|
||||
# [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
|
||||
# [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
|
||||
# num_nodes = 5, num_n1 = 2, num_n2 = 3
|
||||
ntypes = {"n1": 0, "n2": 1}
|
||||
etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
|
||||
indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
|
||||
indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
|
||||
type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
|
||||
node_type_offset = torch.LongTensor([0, 2, 5])
|
||||
return gb.fused_csc_sampling_graph(
|
||||
indptr,
|
||||
indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
)
|
||||
|
||||
|
||||
def test_NegativeSampler_Hetero_Data():
|
||||
graph = get_hetero_graph().to(F.ctx())
|
||||
itemset = gb.HeteroItemSet(
|
||||
{
|
||||
"n1:e1:n2": gb.ItemSet(
|
||||
torch.LongTensor([[0, 0, 1, 1], [0, 2, 0, 1]]).T,
|
||||
names="seeds",
|
||||
),
|
||||
"n2:e2:n1": gb.ItemSet(
|
||||
torch.LongTensor([[0, 0, 1, 1, 2, 2], [0, 1, 1, 0, 0, 1]]).T,
|
||||
names="seeds",
|
||||
),
|
||||
}
|
||||
)
|
||||
batch_size = 2
|
||||
negative_ratio = 1
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=batch_size).copy_to(
|
||||
F.ctx()
|
||||
)
|
||||
negative_dp = gb.UniformNegativeSampler(item_sampler, graph, negative_ratio)
|
||||
assert len(list(negative_dp)) == 5
|
||||
# Perform negative sampling.
|
||||
expected_neg_src = [
|
||||
{"n1:e1:n2": torch.tensor([0, 0])},
|
||||
{"n1:e1:n2": torch.tensor([1, 1])},
|
||||
{"n2:e2:n1": torch.tensor([0, 0])},
|
||||
{"n2:e2:n1": torch.tensor([1, 1])},
|
||||
{"n2:e2:n1": torch.tensor([2, 2])},
|
||||
]
|
||||
for i, data in enumerate(negative_dp):
|
||||
# Check negative seeds value.
|
||||
for etype, seeds_data in data.seeds.items():
|
||||
neg_src = seeds_data[batch_size:, 0]
|
||||
neg_dst = seeds_data[batch_size:, 1]
|
||||
assert torch.equal(expected_neg_src[i][etype].to(F.ctx()), neg_src)
|
||||
assert (neg_dst < 3).all(), neg_dst
|
||||
@@ -0,0 +1,161 @@
|
||||
import unittest
|
||||
from functools import partial
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def get_hetero_graph(include_original_edge_ids):
|
||||
# COO graph:
|
||||
# [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
|
||||
# [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
|
||||
# [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
|
||||
# num_nodes = 5, num_n1 = 2, num_n2 = 3
|
||||
ntypes = {"n1": 0, "n2": 1, "n3": 2}
|
||||
etypes = {"n2:e1:n3": 0, "n3:e2:n2": 1}
|
||||
indptr = torch.LongTensor([0, 0, 2, 4, 6, 8, 10])
|
||||
indices = torch.LongTensor([3, 5, 3, 4, 1, 2, 2, 1, 1, 2])
|
||||
type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
|
||||
edge_attributes = {
|
||||
"weight": torch.FloatTensor(
|
||||
[2.5, 0, 8.4, 0, 0.4, 1.2, 2.5, 0, 8.4, 0.5]
|
||||
),
|
||||
"mask": torch.BoolTensor([1, 0, 1, 0, 1, 1, 1, 0, 1, 1]),
|
||||
}
|
||||
if include_original_edge_ids:
|
||||
edge_attributes[gb.ORIGINAL_EDGE_ID] = (
|
||||
torch.arange(indices.size(0), 0, -1) - 1
|
||||
)
|
||||
node_type_offset = torch.LongTensor([0, 1, 3, 6])
|
||||
return gb.fused_csc_sampling_graph(
|
||||
indptr,
|
||||
indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
edge_attributes=edge_attributes,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(F._default_context_str != "gpu", reason="Enabled only on GPU.")
|
||||
@pytest.mark.parametrize("hetero", [False, True])
|
||||
@pytest.mark.parametrize("prob_name", [None, "weight", "mask"])
|
||||
@pytest.mark.parametrize("sorted", [False, True])
|
||||
@pytest.mark.parametrize("num_cached_edges", [0, 10])
|
||||
@pytest.mark.parametrize("is_pinned", [False, True])
|
||||
@pytest.mark.parametrize("has_orig_edge_ids", [False, True])
|
||||
def test_NeighborSampler_GraphFetch(
|
||||
hetero, prob_name, sorted, num_cached_edges, is_pinned, has_orig_edge_ids
|
||||
):
|
||||
if sorted:
|
||||
items = torch.arange(3)
|
||||
else:
|
||||
items = torch.tensor([2, 0, 1])
|
||||
names = "seeds"
|
||||
itemset = gb.ItemSet(items, names=names)
|
||||
graph = get_hetero_graph(has_orig_edge_ids)
|
||||
graph = graph.pin_memory_() if is_pinned else graph.to(F.ctx())
|
||||
if hetero:
|
||||
itemset = gb.HeteroItemSet({"n3": itemset})
|
||||
else:
|
||||
graph.type_per_edge = None
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
|
||||
fanout = torch.LongTensor([2])
|
||||
preprocess_fn = partial(
|
||||
gb.SubgraphSampler._preprocess, cooperative=False, async_op=False
|
||||
)
|
||||
datapipe = item_sampler.map(preprocess_fn)
|
||||
datapipe = datapipe.map(
|
||||
partial(gb.NeighborSampler._prepare, graph.node_type_to_id)
|
||||
)
|
||||
sample_per_layer = gb.SamplePerLayer(
|
||||
datapipe, graph.sample_neighbors, fanout, False, prob_name, False
|
||||
)
|
||||
compact_per_layer = sample_per_layer.compact_per_layer(True)
|
||||
gb.seed(123)
|
||||
expected_results = list(compact_per_layer)
|
||||
if num_cached_edges > 0:
|
||||
graph._initialize_gpu_graph_cache(num_cached_edges, 1, prob_name)
|
||||
datapipe = datapipe.sample_per_layer(
|
||||
graph.sample_neighbors, fanout, False, prob_name, True
|
||||
)
|
||||
datapipe = datapipe.compact_per_layer(True)
|
||||
gb.seed(123)
|
||||
new_results = list(datapipe)
|
||||
assert len(expected_results) == len(new_results)
|
||||
for a, b in zip(expected_results, new_results):
|
||||
assert repr(a) == repr(b)
|
||||
|
||||
def remove_input_nodes(minibatch):
|
||||
minibatch.input_nodes = None
|
||||
return minibatch
|
||||
|
||||
datapipe = item_sampler.sample_neighbor(
|
||||
graph, [fanout], False, prob_name=prob_name, overlap_fetch=True
|
||||
)
|
||||
datapipe = datapipe.transform(remove_input_nodes)
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
gb.seed(123)
|
||||
new_results = list(dataloader)
|
||||
assert len(expected_results) == len(new_results)
|
||||
for a, b in zip(expected_results, new_results):
|
||||
assert repr(a) == repr(b)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layer_dependency", [False, True])
|
||||
@pytest.mark.parametrize("overlap_graph_fetch", [False, True])
|
||||
def test_labor_dependent_minibatching(layer_dependency, overlap_graph_fetch):
|
||||
if F._default_context_str != "gpu" and overlap_graph_fetch:
|
||||
pytest.skip("overlap_graph_fetch is only available for GPU.")
|
||||
num_edges = 200
|
||||
csc_indptr = torch.cat(
|
||||
(
|
||||
torch.zeros(1, dtype=torch.int64),
|
||||
torch.ones(num_edges + 1, dtype=torch.int64) * num_edges,
|
||||
)
|
||||
)
|
||||
indices = torch.arange(1, num_edges + 1)
|
||||
graph = gb.fused_csc_sampling_graph(
|
||||
csc_indptr.int(),
|
||||
indices.int(),
|
||||
).to(F.ctx())
|
||||
torch.random.set_rng_state(torch.manual_seed(123).get_state())
|
||||
batch_dependency = 100
|
||||
itemset = gb.ItemSet(torch.zeros(batch_dependency + 1).int(), names="seeds")
|
||||
datapipe = gb.ItemSampler(itemset, batch_size=1).copy_to(F.ctx())
|
||||
fanouts = [5, 5]
|
||||
datapipe = datapipe.sample_layer_neighbor(
|
||||
graph,
|
||||
fanouts,
|
||||
overlap_fetch=overlap_graph_fetch,
|
||||
layer_dependency=layer_dependency,
|
||||
batch_dependency=batch_dependency,
|
||||
)
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
res = list(dataloader)
|
||||
assert len(res) == batch_dependency + 1
|
||||
if layer_dependency:
|
||||
assert torch.equal(
|
||||
res[0].input_nodes,
|
||||
res[0].sampled_subgraphs[1].original_row_node_ids,
|
||||
)
|
||||
else:
|
||||
assert res[0].input_nodes.size(0) > res[0].sampled_subgraphs[
|
||||
1
|
||||
].original_row_node_ids.size(0)
|
||||
delta = 0
|
||||
for i in range(batch_dependency):
|
||||
res_current = (
|
||||
res[i].sampled_subgraphs[-1].original_row_node_ids.tolist()
|
||||
)
|
||||
res_next = (
|
||||
res[i + 1].sampled_subgraphs[-1].original_row_node_ids.tolist()
|
||||
)
|
||||
intersect_len = len(set(res_current).intersection(set(res_next)))
|
||||
assert intersect_len >= fanouts[-1]
|
||||
delta += 1 + fanouts[-1] - intersect_len
|
||||
assert delta >= fanouts[-1]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,735 @@
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
from dgl.graphbolt.impl.sampled_subgraph_impl import SampledSubgraphImpl
|
||||
|
||||
|
||||
def _assert_container_equal(lhs, rhs):
|
||||
if isinstance(lhs, torch.Tensor):
|
||||
assert isinstance(rhs, torch.Tensor)
|
||||
assert torch.equal(lhs, rhs)
|
||||
elif isinstance(lhs, tuple):
|
||||
assert isinstance(rhs, tuple)
|
||||
assert len(lhs) == len(rhs)
|
||||
for l, r in zip(lhs, rhs):
|
||||
_assert_container_equal(l, r)
|
||||
elif isinstance(lhs, gb.CSCFormatBase):
|
||||
assert isinstance(rhs, gb.CSCFormatBase)
|
||||
assert len(lhs.indptr) == len(rhs.indptr)
|
||||
assert len(lhs.indices) == len(rhs.indices)
|
||||
_assert_container_equal(lhs.indptr, rhs.indptr)
|
||||
_assert_container_equal(lhs.indices, rhs.indices)
|
||||
elif isinstance(lhs, dict):
|
||||
assert isinstance(rhs, dict)
|
||||
assert len(lhs) == len(rhs)
|
||||
for key, value in lhs.items():
|
||||
assert key in rhs
|
||||
_assert_container_equal(value, rhs[key])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_homo_deduplicated(reverse_row, reverse_column):
|
||||
csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 2, 2, 3]), indices=torch.tensor([0, 3, 2])
|
||||
)
|
||||
if reverse_row:
|
||||
original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
src_to_exclude = torch.tensor([11])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2])
|
||||
|
||||
if reverse_column:
|
||||
original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
dst_to_exclude = torch.tensor([9])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([4])
|
||||
original_edge_ids = torch.Tensor([5, 9, 10])
|
||||
subgraph = SampledSubgraphImpl(
|
||||
csc_formats,
|
||||
original_column_node_ids,
|
||||
original_row_node_ids,
|
||||
original_edge_ids,
|
||||
)
|
||||
edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(2, -1).T
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 2, 2, 2]), indices=torch.tensor([0, 3])
|
||||
)
|
||||
if reverse_row:
|
||||
expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = torch.Tensor([5, 9])
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_homo_duplicated(reverse_row, reverse_column):
|
||||
csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 3, 3, 5]),
|
||||
indices=torch.tensor([0, 3, 3, 2, 2]),
|
||||
)
|
||||
if reverse_row:
|
||||
original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
src_to_exclude = torch.tensor([24])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([3])
|
||||
|
||||
if reverse_column:
|
||||
original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
dst_to_exclude = torch.tensor([11])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([2])
|
||||
original_edge_ids = torch.Tensor([5, 9, 9, 10, 10])
|
||||
subgraph = SampledSubgraphImpl(
|
||||
csc_formats,
|
||||
original_column_node_ids,
|
||||
original_row_node_ids,
|
||||
original_edge_ids,
|
||||
)
|
||||
edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(2, -1).T
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 1, 1, 3]), indices=torch.tensor([0, 2, 2])
|
||||
)
|
||||
if reverse_row:
|
||||
expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = torch.Tensor([5, 10, 10])
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_hetero_deduplicated(reverse_row, reverse_column):
|
||||
csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3]),
|
||||
indices=torch.tensor([2, 1, 0]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
original_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
src_to_exclude = torch.tensor([15, 13])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2, 0])
|
||||
if reverse_column:
|
||||
original_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
dst_to_exclude = torch.tensor([10, 12])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([0, 2])
|
||||
original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
|
||||
subgraph = SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats,
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
|
||||
edges_to_exclude = {
|
||||
"A:relation:B": torch.cat(
|
||||
(
|
||||
src_to_exclude,
|
||||
dst_to_exclude,
|
||||
)
|
||||
)
|
||||
.view(2, -1)
|
||||
.T
|
||||
}
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 1]),
|
||||
indices=torch.tensor([1]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
expected_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = {"A:relation:B": torch.tensor([20])}
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_hetero_duplicated(reverse_row, reverse_column):
|
||||
csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 5]),
|
||||
indices=torch.tensor([2, 2, 1, 1, 0]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
original_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
src_to_exclude = torch.tensor([15, 13])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2, 0])
|
||||
if reverse_column:
|
||||
original_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
dst_to_exclude = torch.tensor([10, 12])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([0, 2])
|
||||
original_edge_ids = {"A:relation:B": torch.tensor([19, 19, 20, 20, 21])}
|
||||
subgraph = SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats,
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
|
||||
edges_to_exclude = {
|
||||
"A:relation:B": torch.cat(
|
||||
(
|
||||
src_to_exclude,
|
||||
dst_to_exclude,
|
||||
)
|
||||
)
|
||||
.view(2, -1)
|
||||
.T
|
||||
}
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 2, 2]),
|
||||
indices=torch.tensor([1, 1]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
expected_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = {"A:relation:B": torch.tensor([20, 20])}
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_homo_deduplicated_tensor(reverse_row, reverse_column):
|
||||
csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 2, 2, 3]), indices=torch.tensor([0, 3, 2])
|
||||
)
|
||||
if reverse_row:
|
||||
original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
src_to_exclude = torch.tensor([11])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2])
|
||||
|
||||
if reverse_column:
|
||||
original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
dst_to_exclude = torch.tensor([9])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([4])
|
||||
original_edge_ids = torch.Tensor([5, 9, 10])
|
||||
subgraph = SampledSubgraphImpl(
|
||||
csc_formats,
|
||||
original_column_node_ids,
|
||||
original_row_node_ids,
|
||||
original_edge_ids,
|
||||
)
|
||||
edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(1, -1)
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 2, 2, 2]), indices=torch.tensor([0, 3])
|
||||
)
|
||||
if reverse_row:
|
||||
expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = torch.Tensor([5, 9])
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_homo_duplicated_tensor(reverse_row, reverse_column):
|
||||
csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 3, 3, 5]),
|
||||
indices=torch.tensor([0, 3, 3, 2, 2]),
|
||||
)
|
||||
if reverse_row:
|
||||
original_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
src_to_exclude = torch.tensor([24])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([3])
|
||||
|
||||
if reverse_column:
|
||||
original_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
dst_to_exclude = torch.tensor([11])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([2])
|
||||
original_edge_ids = torch.Tensor([5, 9, 9, 10, 10])
|
||||
subgraph = SampledSubgraphImpl(
|
||||
csc_formats,
|
||||
original_column_node_ids,
|
||||
original_row_node_ids,
|
||||
original_edge_ids,
|
||||
)
|
||||
edges_to_exclude = torch.cat((src_to_exclude, dst_to_exclude)).view(1, -1)
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 1, 1, 3]), indices=torch.tensor([0, 2, 2])
|
||||
)
|
||||
if reverse_row:
|
||||
expected_row_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = torch.tensor([10, 15, 11, 24, 9])
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = torch.Tensor([5, 10, 10])
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_hetero_deduplicated_tensor(reverse_row, reverse_column):
|
||||
csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3]),
|
||||
indices=torch.tensor([2, 1, 0]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
original_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
src_to_exclude = torch.tensor([15, 13])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2, 0])
|
||||
if reverse_column:
|
||||
original_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
dst_to_exclude = torch.tensor([10, 12])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([0, 2])
|
||||
original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
|
||||
subgraph = SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats,
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
|
||||
edges_to_exclude = {
|
||||
"A:relation:B": torch.cat((src_to_exclude, dst_to_exclude))
|
||||
.view(2, -1)
|
||||
.T
|
||||
}
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 1, 1]),
|
||||
indices=torch.tensor([1]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
expected_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = {"A:relation:B": torch.tensor([20])}
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_row", [True, False])
|
||||
@pytest.mark.parametrize("reverse_column", [True, False])
|
||||
def test_exclude_edges_hetero_duplicated_tensor(reverse_row, reverse_column):
|
||||
csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 5]),
|
||||
indices=torch.tensor([2, 2, 1, 1, 0]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
original_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
src_to_exclude = torch.tensor([15, 13])
|
||||
else:
|
||||
original_row_node_ids = None
|
||||
src_to_exclude = torch.tensor([2, 0])
|
||||
if reverse_column:
|
||||
original_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
dst_to_exclude = torch.tensor([10, 12])
|
||||
else:
|
||||
original_column_node_ids = None
|
||||
dst_to_exclude = torch.tensor([0, 2])
|
||||
original_edge_ids = {"A:relation:B": torch.tensor([19, 19, 20, 20, 21])}
|
||||
subgraph = SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats,
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
|
||||
edges_to_exclude = {
|
||||
"A:relation:B": torch.cat((src_to_exclude, dst_to_exclude))
|
||||
.view(2, -1)
|
||||
.T
|
||||
}
|
||||
result = subgraph.exclude_edges(edges_to_exclude)
|
||||
expected_csc_formats = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 2, 2]),
|
||||
indices=torch.tensor([1, 1]),
|
||||
)
|
||||
}
|
||||
if reverse_row:
|
||||
expected_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
else:
|
||||
expected_row_node_ids = None
|
||||
if reverse_column:
|
||||
expected_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
else:
|
||||
expected_column_node_ids = None
|
||||
expected_edge_ids = {"A:relation:B": torch.tensor([20, 20])}
|
||||
|
||||
_assert_container_equal(result.sampled_csc, expected_csc_formats)
|
||||
_assert_container_equal(
|
||||
result.original_column_node_ids, expected_column_node_ids
|
||||
)
|
||||
_assert_container_equal(result.original_row_node_ids, expected_row_node_ids)
|
||||
_assert_container_equal(result.original_edge_ids, expected_edge_ids)
|
||||
|
||||
|
||||
def test_to_pyg_homo():
|
||||
graph = dgl.graph(([5, 0, 7, 7, 2, 4], [0, 1, 2, 2, 3, 4]))
|
||||
graph = gb.from_dglgraph(graph, is_homogeneous=True).to(F.ctx())
|
||||
items = torch.LongTensor([[0, 3], [4, 4]])
|
||||
names = "seeds"
|
||||
itemset = gb.ItemSet(items, names=names)
|
||||
datapipe = gb.ItemSampler(itemset, batch_size=4).copy_to(F.ctx())
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([-1]) for _ in range(num_layer)]
|
||||
sampler = gb.NeighborSampler
|
||||
datapipe = sampler(
|
||||
datapipe,
|
||||
graph,
|
||||
fanouts,
|
||||
deduplicate=True,
|
||||
)
|
||||
for minibatch in datapipe:
|
||||
x = torch.randn((minibatch.node_ids().size(0), 2), dtype=torch.float32)
|
||||
for subgraph in minibatch.sampled_subgraphs:
|
||||
(x_src, x_dst), edge_index, sizes = subgraph.to_pyg(x)
|
||||
assert torch.equal(x_src, x)
|
||||
dst_size = subgraph.original_column_node_ids.size(0)
|
||||
assert torch.equal(x_dst, x[:dst_size])
|
||||
src_size = subgraph.original_row_node_ids.size(0)
|
||||
assert dst_size == sizes[1]
|
||||
assert src_size == sizes[0]
|
||||
assert torch.equal(edge_index[0], subgraph.sampled_csc.indices)
|
||||
assert torch.equal(
|
||||
edge_index[1],
|
||||
gb.expand_indptr(
|
||||
subgraph.sampled_csc.indptr,
|
||||
subgraph.sampled_csc.indices.dtype,
|
||||
),
|
||||
)
|
||||
x = x_dst
|
||||
|
||||
|
||||
def test_to_pyg_hetero():
|
||||
# COO graph:
|
||||
# [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
|
||||
# [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
|
||||
# [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
|
||||
# num_nodes = 5, num_n1 = 2, num_n2 = 3
|
||||
ntypes = {"n1": 0, "n2": 1}
|
||||
etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
|
||||
indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
|
||||
indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
|
||||
type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
|
||||
node_type_offset = torch.LongTensor([0, 2, 5])
|
||||
graph = gb.fused_csc_sampling_graph(
|
||||
indptr,
|
||||
indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
).to(F.ctx())
|
||||
itemset = gb.HeteroItemSet(
|
||||
{"n1:e1:n2": gb.ItemSet(torch.tensor([[0, 1]]), names="seeds")}
|
||||
)
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
Sampler = gb.NeighborSampler
|
||||
datapipe = Sampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
fanouts,
|
||||
deduplicate=True,
|
||||
)
|
||||
for minibatch in datapipe:
|
||||
x = {}
|
||||
for key, ids in minibatch.node_ids().items():
|
||||
x[key] = torch.randn((ids.size(0), 2), dtype=torch.float32)
|
||||
for subgraph in minibatch.sampled_subgraphs:
|
||||
(x_src, x_dst), edge_index, sizes = subgraph.to_pyg(x)
|
||||
assert x_src == x
|
||||
for ntype in x:
|
||||
dst_size = subgraph.original_column_node_ids[ntype].size(0)
|
||||
assert torch.equal(x_dst[ntype], x[ntype][:dst_size])
|
||||
for etype in subgraph.sampled_csc:
|
||||
src_ntype, _, dst_ntype = gb.etype_str_to_tuple(etype)
|
||||
src_size = subgraph.original_row_node_ids[src_ntype].size(0)
|
||||
dst_size = subgraph.original_column_node_ids[dst_ntype].size(0)
|
||||
assert dst_size == sizes[etype][1]
|
||||
assert src_size == sizes[etype][0]
|
||||
assert torch.equal(
|
||||
edge_index[etype][0], subgraph.sampled_csc[etype].indices
|
||||
)
|
||||
assert torch.equal(
|
||||
edge_index[etype][1],
|
||||
gb.expand_indptr(
|
||||
subgraph.sampled_csc[etype].indptr,
|
||||
subgraph.sampled_csc[etype].indices.dtype,
|
||||
),
|
||||
)
|
||||
x = x_dst
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "cpu",
|
||||
reason="`to` function needs GPU to test.",
|
||||
)
|
||||
def test_sampled_subgraph_to_device():
|
||||
# Initialize data.
|
||||
csc_format = {
|
||||
"A:relation:B": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3]),
|
||||
indices=torch.tensor([0, 1, 2]),
|
||||
)
|
||||
}
|
||||
original_row_node_ids = {
|
||||
"A": torch.tensor([13, 14, 15]),
|
||||
}
|
||||
src_to_exclude = torch.tensor([15, 13])
|
||||
original_column_node_ids = {
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
}
|
||||
dst_to_exclude = torch.tensor([10, 12])
|
||||
original_edge_ids = {"A:relation:B": torch.tensor([19, 20, 21])}
|
||||
subgraph = SampledSubgraphImpl(
|
||||
sampled_csc=csc_format,
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
edges_to_exclude = {
|
||||
"A:relation:B": torch.cat(
|
||||
(
|
||||
src_to_exclude,
|
||||
dst_to_exclude,
|
||||
)
|
||||
)
|
||||
.view(2, -1)
|
||||
.T
|
||||
}
|
||||
graph = subgraph.exclude_edges(edges_to_exclude)
|
||||
|
||||
# Copy to device.
|
||||
graph = graph.to("cuda")
|
||||
|
||||
# Check.
|
||||
for key in graph.sampled_csc:
|
||||
assert graph.sampled_csc[key].indices.device.type == "cuda"
|
||||
assert graph.sampled_csc[key].indptr.device.type == "cuda"
|
||||
for key in graph.original_column_node_ids:
|
||||
assert graph.original_column_node_ids[key].device.type == "cuda"
|
||||
for key in graph.original_row_node_ids:
|
||||
assert graph.original_row_node_ids[key].device.type == "cuda"
|
||||
for key in graph.original_edge_ids:
|
||||
assert graph.original_edge_ids[key].device.type == "cuda"
|
||||
|
||||
|
||||
def test_sampled_subgraph_impl_representation_homo():
|
||||
sampled_subgraph_impl = SampledSubgraphImpl(
|
||||
sampled_csc=gb.CSCFormatBase(
|
||||
indptr=torch.arange(0, 101, 10),
|
||||
indices=torch.arange(10, 110),
|
||||
),
|
||||
original_column_node_ids=torch.arange(0, 10),
|
||||
original_row_node_ids=torch.arange(0, 110),
|
||||
original_edge_ids=None,
|
||||
)
|
||||
expected_result = str(
|
||||
"""SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([ 0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]),
|
||||
indices=tensor([ 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23,
|
||||
24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37,
|
||||
38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51,
|
||||
52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65,
|
||||
66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
|
||||
80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93,
|
||||
94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107,
|
||||
108, 109]),
|
||||
),
|
||||
original_row_node_ids=tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,
|
||||
14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,
|
||||
28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41,
|
||||
42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55,
|
||||
56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69,
|
||||
70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83,
|
||||
84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97,
|
||||
98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109]),
|
||||
original_edge_ids=None,
|
||||
original_column_node_ids=tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
|
||||
)"""
|
||||
)
|
||||
assert str(sampled_subgraph_impl) == expected_result, print(
|
||||
sampled_subgraph_impl
|
||||
)
|
||||
|
||||
|
||||
def test_sampled_subgraph_impl_representation_hetero():
|
||||
sampled_subgraph_impl = SampledSubgraphImpl(
|
||||
sampled_csc={
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4]),
|
||||
indices=torch.tensor([4, 5, 6, 7]),
|
||||
),
|
||||
"n2:e2:n1": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 6, 8]),
|
||||
indices=torch.tensor([2, 3, 4, 5, 6, 7, 8, 9]),
|
||||
),
|
||||
},
|
||||
original_column_node_ids={
|
||||
"n1": torch.tensor([1, 0, 0, 1]),
|
||||
"n2": torch.tensor([1, 2]),
|
||||
},
|
||||
original_row_node_ids={
|
||||
"n1": torch.tensor([1, 0, 0, 1, 1, 0, 0, 1]),
|
||||
"n2": torch.tensor([1, 2, 0, 1, 0, 2, 0, 2, 0, 1]),
|
||||
},
|
||||
original_edge_ids=None,
|
||||
)
|
||||
expected_result = str(
|
||||
"""SampledSubgraphImpl(sampled_csc={'n1:e1:n2': CSCFormatBase(indptr=tensor([0, 2, 4]),
|
||||
indices=tensor([4, 5, 6, 7]),
|
||||
), 'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 2, 4, 6, 8]),
|
||||
indices=tensor([2, 3, 4, 5, 6, 7, 8, 9]),
|
||||
)},
|
||||
original_row_node_ids={'n1': tensor([1, 0, 0, 1, 1, 0, 0, 1]), 'n2': tensor([1, 2, 0, 1, 0, 2, 0, 2, 0, 1])},
|
||||
original_edge_ids=None,
|
||||
original_column_node_ids={'n1': tensor([1, 0, 0, 1]), 'n2': tensor([1, 2])},
|
||||
)"""
|
||||
)
|
||||
assert str(sampled_subgraph_impl) == expected_result, print(
|
||||
sampled_subgraph_impl
|
||||
)
|
||||
@@ -0,0 +1,458 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import backend as F
|
||||
|
||||
import numpy as np
|
||||
import pydantic
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def to_on_disk_tensor(test_dir, name, t):
|
||||
path = os.path.join(test_dir, name + ".npy")
|
||||
t = t.numpy()
|
||||
np.save(path, t)
|
||||
# The Pytorch tensor is a view of the numpy array on disk, which does not
|
||||
# consume memory.
|
||||
t = torch.as_tensor(np.load(path, mmap_mode="r+"))
|
||||
return t
|
||||
|
||||
|
||||
@pytest.mark.parametrize("in_memory", [True, False])
|
||||
def test_torch_based_feature(in_memory):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
metadata = {"max_value": 3}
|
||||
if not in_memory:
|
||||
a = to_on_disk_tensor(test_dir, "a", a)
|
||||
b = to_on_disk_tensor(test_dir, "b", b)
|
||||
|
||||
feature_a = gb.TorchBasedFeature(a, metadata=metadata)
|
||||
feature_b = gb.TorchBasedFeature(b)
|
||||
|
||||
# Read the entire feature.
|
||||
assert torch.equal(
|
||||
feature_a.read(), torch.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
)
|
||||
|
||||
# Test read the feature with ids.
|
||||
assert torch.equal(
|
||||
feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
)
|
||||
# Read the feature with ids.
|
||||
assert torch.equal(
|
||||
feature_a.read(torch.tensor([0])),
|
||||
torch.tensor([[1, 2, 3]]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_b.read(torch.tensor([1])),
|
||||
torch.tensor([[[4, 5], [6, 7]]]),
|
||||
)
|
||||
# Update the feature with ids.
|
||||
feature_a.update(torch.tensor([[0, 1, 2]]), torch.tensor([0]))
|
||||
assert torch.equal(
|
||||
feature_a.read(), torch.tensor([[0, 1, 2], [4, 5, 6]])
|
||||
)
|
||||
feature_b.update(torch.tensor([[[1, 2], [3, 4]]]), torch.tensor([1]))
|
||||
assert torch.equal(
|
||||
feature_b.read(), torch.tensor([[[1, 2], [3, 4]], [[1, 2], [3, 4]]])
|
||||
)
|
||||
|
||||
# Test update the feature.
|
||||
feature_a.update(torch.tensor([[5, 1, 3]]))
|
||||
assert torch.equal(
|
||||
feature_a.read(),
|
||||
torch.tensor([[5, 1, 3]]),
|
||||
), print(feature_a.read())
|
||||
feature_b.update(
|
||||
torch.tensor([[[1, 3], [5, 7]], [[2, 4], [6, 8]], [[2, 4], [6, 8]]])
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_b.read(),
|
||||
torch.tensor(
|
||||
[[[1, 3], [5, 7]], [[2, 4], [6, 8]], [[2, 4], [6, 8]]]
|
||||
),
|
||||
)
|
||||
|
||||
# Test get the size and count of the entire feature.
|
||||
assert feature_a.size() == torch.Size([3])
|
||||
assert feature_b.size() == torch.Size([2, 2])
|
||||
assert feature_a.count() == 1
|
||||
assert feature_b.count() == 3
|
||||
|
||||
# Test get metadata of the feature.
|
||||
assert feature_a.metadata() == metadata
|
||||
assert feature_b.metadata() == {}
|
||||
|
||||
with pytest.raises(IndexError):
|
||||
feature_a.read(torch.tensor([0, 1, 2, 3]))
|
||||
|
||||
# For windows, the file is locked by the numpy.load. We need to delete
|
||||
# it before closing the temporary directory.
|
||||
a = b = None
|
||||
feature_a = feature_b = None
|
||||
|
||||
# Test loaded tensors' contiguity from C/Fortran contiguous ndarray.
|
||||
contiguous_numpy = np.array([[1, 2, 3], [4, 5, 6]], order="C")
|
||||
non_contiguous_numpy = np.array([[1, 2, 3], [4, 5, 6]], order="F")
|
||||
assert contiguous_numpy.flags["C_CONTIGUOUS"]
|
||||
assert non_contiguous_numpy.flags["F_CONTIGUOUS"]
|
||||
np.save(
|
||||
os.path.join(test_dir, "contiguous_numpy.npy"), contiguous_numpy
|
||||
)
|
||||
np.save(
|
||||
os.path.join(test_dir, "non_contiguous_numpy.npy"),
|
||||
non_contiguous_numpy,
|
||||
)
|
||||
|
||||
cur_mmap_mode = None
|
||||
if not in_memory:
|
||||
cur_mmap_mode = "r+"
|
||||
feature_a = gb.TorchBasedFeature(
|
||||
torch.from_numpy(
|
||||
np.load(
|
||||
os.path.join(test_dir, "contiguous_numpy.npy"),
|
||||
mmap_mode=cur_mmap_mode,
|
||||
)
|
||||
)
|
||||
)
|
||||
feature_b = gb.TorchBasedFeature(
|
||||
torch.from_numpy(
|
||||
np.load(
|
||||
os.path.join(test_dir, "non_contiguous_numpy.npy"),
|
||||
mmap_mode=cur_mmap_mode,
|
||||
)
|
||||
)
|
||||
)
|
||||
assert feature_a._tensor.is_contiguous()
|
||||
assert feature_b._tensor.is_contiguous()
|
||||
|
||||
contiguous_numpy = non_contiguous_numpy = None
|
||||
feature_a = feature_b = None
|
||||
|
||||
|
||||
def is_feature_store_on_cuda(store):
|
||||
for feature in store._features.values():
|
||||
assert feature._tensor.is_cuda
|
||||
|
||||
|
||||
def is_feature_store_on_cpu(store):
|
||||
for feature in store._features.values():
|
||||
assert not feature._tensor.is_cuda
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "cpu",
|
||||
reason="Tests for pinned memory are only meaningful on GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize("device", ["pinned", "cuda"])
|
||||
def test_feature_store_to_device(device):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
|
||||
write_tensor_to_disk(test_dir, "a", a, fmt="torch")
|
||||
write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
|
||||
feature_data = [
|
||||
gb.OnDiskFeatureData(
|
||||
domain="node",
|
||||
type="paper",
|
||||
name="a",
|
||||
format="torch",
|
||||
path=os.path.join(test_dir, "a.pt"),
|
||||
),
|
||||
gb.OnDiskFeatureData(
|
||||
domain="edge",
|
||||
type="paper:cites:paper",
|
||||
name="b",
|
||||
format="numpy",
|
||||
path=os.path.join(test_dir, "b.npy"),
|
||||
),
|
||||
]
|
||||
feature_store = gb.TorchBasedFeatureStore(feature_data)
|
||||
feature_store2 = feature_store.to(device)
|
||||
if device == "pinned":
|
||||
assert feature_store2.is_pinned()
|
||||
elif device == "cuda":
|
||||
is_feature_store_on_cuda(feature_store2)
|
||||
|
||||
# The original variable should be untouched.
|
||||
is_feature_store_on_cpu(feature_store)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "cpu",
|
||||
reason="Tests for pinned memory are only meaningful on GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.int8,
|
||||
torch.float16,
|
||||
torch.complex128,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("shape", [(2, 1), (2, 3), (2, 2, 2), (137, 13, 3)])
|
||||
@pytest.mark.parametrize("in_place", [False, True])
|
||||
def test_torch_based_pinned_feature(dtype, idtype, shape, in_place):
|
||||
if dtype == torch.complex128:
|
||||
tensor = torch.complex(
|
||||
torch.randint(0, 13, shape, dtype=torch.float64),
|
||||
torch.randint(0, 13, shape, dtype=torch.float64),
|
||||
)
|
||||
else:
|
||||
tensor = torch.randint(0, 13, shape, dtype=dtype)
|
||||
test_tensor = tensor.clone().detach()
|
||||
test_tensor_cuda = test_tensor.cuda()
|
||||
|
||||
feature = gb.TorchBasedFeature(tensor)
|
||||
if in_place:
|
||||
if gb.is_wsl():
|
||||
pytest.skip("In place pinning is not supported on WSL.")
|
||||
feature.pin_memory_()
|
||||
|
||||
# Check if pinning is truly in-place.
|
||||
assert feature._tensor.data_ptr() == tensor.data_ptr()
|
||||
else:
|
||||
feature = feature.to("pinned")
|
||||
|
||||
assert feature.is_pinned()
|
||||
|
||||
# Test read entire pinned feature, the result should be on cuda.
|
||||
assert torch.equal(feature.read(), test_tensor_cuda)
|
||||
assert feature.read().is_cuda
|
||||
assert torch.equal(
|
||||
feature.read(torch.tensor([0], dtype=idtype).cuda()),
|
||||
test_tensor_cuda[[0]],
|
||||
)
|
||||
|
||||
# Test read pinned feature with idx on cuda, the result should be on cuda.
|
||||
assert feature.read(torch.tensor([0], dtype=idtype).cuda()).is_cuda
|
||||
|
||||
# Test read pinned feature with idx on cpu, the result should be on cpu.
|
||||
assert torch.equal(
|
||||
feature.read(torch.tensor([0], dtype=idtype)), test_tensor[[0]]
|
||||
)
|
||||
assert not feature.read(torch.tensor([0], dtype=idtype)).is_cuda
|
||||
|
||||
|
||||
def write_tensor_to_disk(dir, name, t, fmt="torch"):
|
||||
if fmt == "torch":
|
||||
torch.save(t, os.path.join(dir, name + ".pt"))
|
||||
elif fmt == "numpy":
|
||||
t = t.numpy()
|
||||
np.save(os.path.join(dir, name + ".npy"), t)
|
||||
else:
|
||||
raise ValueError(f"Unsupported format: {fmt}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("in_memory", [True, False])
|
||||
def test_torch_based_feature_store(in_memory):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
|
||||
write_tensor_to_disk(test_dir, "a", a, fmt="torch")
|
||||
write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
|
||||
feature_data = [
|
||||
gb.OnDiskFeatureData(
|
||||
domain="node",
|
||||
type="paper",
|
||||
name="a",
|
||||
format="torch",
|
||||
path=os.path.join(test_dir, "a.pt"),
|
||||
in_memory=True,
|
||||
),
|
||||
gb.OnDiskFeatureData(
|
||||
domain="edge",
|
||||
type="paper:cites:paper",
|
||||
name="b",
|
||||
format="numpy",
|
||||
path=os.path.join(test_dir, "b.npy"),
|
||||
in_memory=in_memory,
|
||||
),
|
||||
]
|
||||
feature_store = gb.TorchBasedFeatureStore(feature_data)
|
||||
|
||||
assert isinstance(
|
||||
feature_store[("node", "paper", "a")], gb.TorchBasedFeature
|
||||
)
|
||||
assert isinstance(
|
||||
feature_store[("edge", "paper:cites:paper", "b")],
|
||||
gb.TorchBasedFeature if in_memory else gb.DiskBasedFeature,
|
||||
)
|
||||
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", "paper", "a"),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
assert torch.equal(
|
||||
feature_store.read("edge", "paper:cites:paper", "b"),
|
||||
torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]]),
|
||||
)
|
||||
|
||||
# Test get the size of the entire feature.
|
||||
assert feature_store.size("node", "paper", "a") == torch.Size([3])
|
||||
assert feature_store.size(
|
||||
"edge", "paper:cites:paper", "b"
|
||||
) == torch.Size([2, 2])
|
||||
|
||||
# Test get the keys of the features.
|
||||
assert feature_store.keys() == [
|
||||
("node", "paper", "a"),
|
||||
("edge", "paper:cites:paper", "b"),
|
||||
]
|
||||
|
||||
# For windows, the file is locked by the numpy.load. We need to delete
|
||||
# it before closing the temporary directory.
|
||||
a = b = None
|
||||
feature_store = None
|
||||
|
||||
# ``domain`` should be enum.
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
_ = gb.OnDiskFeatureData(
|
||||
domain="invalid",
|
||||
type="paper",
|
||||
name="a",
|
||||
format="torch",
|
||||
path=os.path.join(test_dir, "a.pt"),
|
||||
in_memory=True,
|
||||
)
|
||||
|
||||
# ``type`` could be null.
|
||||
feature_data = [
|
||||
gb.OnDiskFeatureData(
|
||||
domain="node",
|
||||
name="a",
|
||||
format="torch",
|
||||
path=os.path.join(test_dir, "a.pt"),
|
||||
in_memory=True,
|
||||
),
|
||||
]
|
||||
feature_store = gb.TorchBasedFeatureStore(feature_data)
|
||||
# Test read the entire feature.
|
||||
assert torch.equal(
|
||||
feature_store.read("node", None, "a"),
|
||||
torch.tensor([[1, 2, 4], [2, 5, 3]]),
|
||||
)
|
||||
# Test get the size of the entire feature.
|
||||
assert feature_store.size("node", None, "a") == torch.Size([3])
|
||||
|
||||
feature_store = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("in_memory", [True, False])
|
||||
def test_torch_based_feature_repr(in_memory):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[4, 5], [6, 7]]])
|
||||
metadata = {"max_value": 3}
|
||||
if not in_memory:
|
||||
a = to_on_disk_tensor(test_dir, "a", a)
|
||||
b = to_on_disk_tensor(test_dir, "b", b)
|
||||
|
||||
feature_a = gb.TorchBasedFeature(a, metadata=metadata)
|
||||
feature_b = gb.TorchBasedFeature(b)
|
||||
|
||||
expected_str_feature_a = (
|
||||
"TorchBasedFeature(\n"
|
||||
" feature=tensor([[1, 2, 3],\n"
|
||||
" [4, 5, 6]]),\n"
|
||||
" metadata={'max_value': 3},\n"
|
||||
")"
|
||||
)
|
||||
expected_str_feature_b = (
|
||||
"TorchBasedFeature(\n"
|
||||
" feature=tensor([[[1, 2],\n"
|
||||
" [3, 4]],\n"
|
||||
"\n"
|
||||
" [[4, 5],\n"
|
||||
" [6, 7]]]),\n"
|
||||
" metadata={},\n"
|
||||
")"
|
||||
)
|
||||
|
||||
assert repr(feature_a) == expected_str_feature_a, feature_a
|
||||
assert repr(feature_b) == expected_str_feature_b, feature_b
|
||||
|
||||
a = b = metadata = None
|
||||
feature_a = feature_b = None
|
||||
expected_str_feature_a = expected_str_feature_b = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("in_memory", [True, False])
|
||||
def test_torch_based_feature_store_repr(in_memory):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
a = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
b = torch.tensor([[[1, 2], [3, 4]], [[2, 5], [3, 4]]])
|
||||
write_tensor_to_disk(test_dir, "a", a, fmt="torch")
|
||||
write_tensor_to_disk(test_dir, "b", b, fmt="numpy")
|
||||
feature_data = [
|
||||
gb.OnDiskFeatureData(
|
||||
domain="node",
|
||||
type="paper",
|
||||
name="a",
|
||||
format="torch",
|
||||
path=os.path.join(test_dir, "a.pt"),
|
||||
in_memory=True,
|
||||
),
|
||||
gb.OnDiskFeatureData(
|
||||
domain="edge",
|
||||
type="paper:cites:paper",
|
||||
name="b",
|
||||
format="numpy",
|
||||
path=os.path.join(test_dir, "b.npy"),
|
||||
in_memory=in_memory,
|
||||
),
|
||||
]
|
||||
feature_store = gb.TorchBasedFeatureStore(feature_data)
|
||||
|
||||
expected_feature_store_str = (
|
||||
(
|
||||
"TorchBasedFeatureStore(\n"
|
||||
" {(<OnDiskFeatureDataDomain.NODE: 'node'>, 'paper', 'a'): TorchBasedFeature(\n"
|
||||
" feature=tensor([[1, 2, 4],\n"
|
||||
" [2, 5, 3]]),\n"
|
||||
" metadata={},\n"
|
||||
" ), (<OnDiskFeatureDataDomain.EDGE: 'edge'>, 'paper:cites:paper', 'b'): TorchBasedFeature(\n"
|
||||
" feature=tensor([[[1, 2],\n"
|
||||
" [3, 4]],\n"
|
||||
"\n"
|
||||
" [[2, 5],\n"
|
||||
" [3, 4]]]),\n"
|
||||
" metadata={},\n"
|
||||
" )}\n"
|
||||
")"
|
||||
)
|
||||
if in_memory
|
||||
else (
|
||||
"TorchBasedFeatureStore(\n"
|
||||
" {(<OnDiskFeatureDataDomain.NODE: 'node'>, 'paper', 'a'): TorchBasedFeature(\n"
|
||||
" feature=tensor([[1, 2, 4],\n"
|
||||
" [2, 5, 3]]),\n"
|
||||
" metadata={},\n"
|
||||
" ), (<OnDiskFeatureDataDomain.EDGE: 'edge'>, 'paper:cites:paper', 'b'): DiskBasedFeature(\n"
|
||||
" feature=tensor([[[1, 2],\n"
|
||||
" [3, 4]],\n"
|
||||
"\n"
|
||||
" [[2, 5],\n"
|
||||
" [3, 4]]]),\n"
|
||||
" metadata={},\n"
|
||||
" )}\n"
|
||||
")"
|
||||
)
|
||||
)
|
||||
|
||||
assert repr(feature_store) == expected_feature_store_str, feature_store
|
||||
|
||||
a = b = feature_data = None
|
||||
feature_store = expected_feature_store_str = None
|
||||
@@ -0,0 +1,273 @@
|
||||
import backend as F
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def test_unique_and_compact_hetero():
|
||||
N1 = torch.tensor(
|
||||
[0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
|
||||
)
|
||||
N2 = torch.tensor([0, 3, 3, 5, 2, 7, 2, 8, 4, 9, 2, 3], device=F.ctx())
|
||||
N3 = torch.tensor([1, 2, 6, 6, 1, 8, 3, 6, 3, 2], device=F.ctx())
|
||||
expected_unique = {
|
||||
"n1": torch.tensor([0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()),
|
||||
"n2": torch.tensor([0, 3, 5, 2, 7, 8, 4, 9], device=F.ctx()),
|
||||
"n3": torch.tensor([1, 2, 6, 8, 3], device=F.ctx()),
|
||||
}
|
||||
if N1.is_cuda and torch.cuda.get_device_capability()[0] < 7:
|
||||
expected_reverse_id = {
|
||||
k: v.sort()[1] for k, v in expected_unique.items()
|
||||
}
|
||||
expected_unique = {k: v.sort()[0] for k, v in expected_unique.items()}
|
||||
else:
|
||||
expected_reverse_id = {
|
||||
k: torch.arange(0, v.shape[0], device=F.ctx())
|
||||
for k, v in expected_unique.items()
|
||||
}
|
||||
nodes_dict = {
|
||||
"n1": N1.split(5),
|
||||
"n2": N2.split(4),
|
||||
"n3": N3.split(2),
|
||||
}
|
||||
expected_nodes_dict = {
|
||||
"n1": [
|
||||
torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
|
||||
torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
|
||||
torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
|
||||
],
|
||||
"n2": [
|
||||
torch.tensor([0, 1, 1, 2], device=F.ctx()),
|
||||
torch.tensor([3, 4, 3, 5], device=F.ctx()),
|
||||
torch.tensor([6, 7, 3, 1], device=F.ctx()),
|
||||
],
|
||||
"n3": [
|
||||
torch.tensor([0, 1], device=F.ctx()),
|
||||
torch.tensor([2, 2], device=F.ctx()),
|
||||
torch.tensor([0, 3], device=F.ctx()),
|
||||
torch.tensor([4, 2], device=F.ctx()),
|
||||
torch.tensor([4, 1], device=F.ctx()),
|
||||
],
|
||||
}
|
||||
|
||||
unique, compacted, _ = gb.unique_and_compact(nodes_dict)
|
||||
for ntype, nodes in unique.items():
|
||||
expected_nodes = expected_unique[ntype]
|
||||
assert torch.equal(nodes, expected_nodes)
|
||||
|
||||
for ntype, nodes in compacted.items():
|
||||
expected_nodes = expected_nodes_dict[ntype]
|
||||
assert isinstance(nodes, list)
|
||||
for expected_node, node in zip(expected_nodes, nodes):
|
||||
node = expected_reverse_id[ntype][node]
|
||||
assert torch.equal(expected_node, node)
|
||||
|
||||
|
||||
def test_unique_and_compact_homo():
|
||||
N = torch.tensor(
|
||||
[0, 5, 2, 7, 12, 7, 9, 5, 6, 2, 3, 4, 1, 0, 9], device=F.ctx()
|
||||
)
|
||||
expected_unique_N = torch.tensor(
|
||||
[0, 5, 2, 7, 12, 9, 6, 3, 4, 1], device=F.ctx()
|
||||
)
|
||||
if N.is_cuda and torch.cuda.get_device_capability()[0] < 7:
|
||||
expected_reverse_id_N = expected_unique_N.sort()[1]
|
||||
expected_unique_N = expected_unique_N.sort()[0]
|
||||
else:
|
||||
expected_reverse_id_N = torch.arange(
|
||||
0, expected_unique_N.shape[0], device=F.ctx()
|
||||
)
|
||||
nodes_list = N.split(5)
|
||||
expected_nodes_list = [
|
||||
torch.tensor([0, 1, 2, 3, 4], device=F.ctx()),
|
||||
torch.tensor([3, 5, 1, 6, 2], device=F.ctx()),
|
||||
torch.tensor([7, 8, 9, 0, 5], device=F.ctx()),
|
||||
]
|
||||
|
||||
unique, compacted, _ = gb.unique_and_compact(nodes_list)
|
||||
|
||||
assert torch.equal(unique, expected_unique_N)
|
||||
assert isinstance(compacted, list)
|
||||
for expected_node, node in zip(expected_nodes_list, compacted):
|
||||
node = expected_reverse_id_N[node]
|
||||
assert torch.equal(expected_node, node)
|
||||
|
||||
|
||||
def test_unique_and_compact_csc_formats_hetero():
|
||||
dst_nodes = {
|
||||
"n2": torch.tensor([2, 4, 1, 3]),
|
||||
"n3": torch.tensor([1, 3, 2, 7]),
|
||||
}
|
||||
csc_formats = {
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 3, 4, 7, 10]),
|
||||
indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
|
||||
),
|
||||
"n1:e2:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 4, 7, 10]),
|
||||
indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
|
||||
),
|
||||
"n2:e3:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 6, 8]),
|
||||
indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
|
||||
),
|
||||
}
|
||||
|
||||
expected_unique_nodes = {
|
||||
"n1": torch.tensor([1, 3, 4, 6, 2, 7, 9, 5, 8, 0]),
|
||||
"n2": torch.tensor([2, 4, 1, 3, 5, 6, 0]),
|
||||
"n3": torch.tensor([1, 3, 2, 7]),
|
||||
}
|
||||
expected_csc_formats = {
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 3, 4, 7, 10]),
|
||||
indices=torch.tensor([0, 1, 2, 3, 4, 5, 6, 2, 4, 3]),
|
||||
),
|
||||
"n1:e2:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 4, 7, 10]),
|
||||
indices=torch.tensor([7, 4, 3, 2, 5, 4, 8, 0, 1, 9]),
|
||||
),
|
||||
"n2:e3:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 6, 8]),
|
||||
indices=torch.tensor([0, 4, 1, 2, 1, 3, 5, 6]),
|
||||
),
|
||||
}
|
||||
|
||||
unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
|
||||
csc_formats, dst_nodes
|
||||
)
|
||||
|
||||
for ntype, nodes in unique_nodes.items():
|
||||
expected_nodes = expected_unique_nodes[ntype]
|
||||
assert torch.equal(nodes, expected_nodes)
|
||||
for etype, pair in compacted_csc_formats.items():
|
||||
indices = pair.indices
|
||||
indptr = pair.indptr
|
||||
expected_indices = expected_csc_formats[etype].indices
|
||||
expected_indptr = expected_csc_formats[etype].indptr
|
||||
assert torch.equal(indices, expected_indices)
|
||||
assert torch.equal(indptr, expected_indptr)
|
||||
|
||||
|
||||
def test_unique_and_compact_csc_formats_homo():
|
||||
seeds = torch.tensor([1, 3, 5, 2, 6])
|
||||
indptr = torch.tensor([0, 2, 4, 6, 7, 11])
|
||||
indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
|
||||
csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
|
||||
|
||||
expected_unique_nodes = torch.tensor([1, 3, 5, 2, 6, 4])
|
||||
expected_indptr = indptr
|
||||
expected_indices = torch.tensor([3, 1, 0, 5, 2, 3, 2, 0, 5, 5, 4])
|
||||
|
||||
unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
|
||||
csc_formats, seeds
|
||||
)
|
||||
|
||||
indptr = compacted_csc_formats.indptr
|
||||
indices = compacted_csc_formats.indices
|
||||
assert torch.equal(indptr, expected_indptr)
|
||||
assert torch.equal(indices, expected_indices)
|
||||
assert torch.equal(unique_nodes, expected_unique_nodes)
|
||||
|
||||
|
||||
def test_unique_and_compact_incorrect_indptr():
|
||||
seeds = torch.tensor([1, 3, 5, 2, 6, 7])
|
||||
indptr = torch.tensor([0, 2, 4, 6, 7, 11])
|
||||
indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
|
||||
csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
|
||||
|
||||
# The number of seeds is not corresponding to indptr.
|
||||
with pytest.raises(AssertionError):
|
||||
gb.unique_and_compact_csc_formats(csc_formats, seeds)
|
||||
|
||||
|
||||
def test_compact_csc_format_hetero():
|
||||
dst_nodes = {
|
||||
"n2": torch.tensor([2, 4, 1, 3]),
|
||||
"n3": torch.tensor([1, 3, 2, 7]),
|
||||
}
|
||||
csc_formats = {
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 3, 4, 7, 10]),
|
||||
indices=torch.tensor([1, 3, 4, 6, 2, 7, 9, 4, 2, 6]),
|
||||
),
|
||||
"n1:e2:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 4, 7, 10]),
|
||||
indices=torch.tensor([5, 2, 6, 4, 7, 2, 8, 1, 3, 0]),
|
||||
),
|
||||
"n2:e3:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 6, 8]),
|
||||
indices=torch.tensor([2, 5, 4, 1, 4, 3, 6, 0]),
|
||||
),
|
||||
}
|
||||
|
||||
expected_original_row_ids = {
|
||||
"n1": torch.tensor(
|
||||
[1, 3, 4, 6, 2, 7, 9, 4, 2, 6, 5, 2, 6, 4, 7, 2, 8, 1, 3, 0]
|
||||
),
|
||||
"n2": torch.tensor([2, 4, 1, 3, 2, 5, 4, 1, 4, 3, 6, 0]),
|
||||
"n3": torch.tensor([1, 3, 2, 7]),
|
||||
}
|
||||
expected_csc_formats = {
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 3, 4, 7, 10]),
|
||||
indices=torch.arange(0, 10),
|
||||
),
|
||||
"n1:e2:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 4, 7, 10]),
|
||||
indices=torch.arange(0, 10) + 10,
|
||||
),
|
||||
"n2:e3:n3": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4, 6, 8]),
|
||||
indices=torch.arange(0, 8) + 4,
|
||||
),
|
||||
}
|
||||
original_row_ids, compacted_csc_formats = gb.compact_csc_format(
|
||||
csc_formats, dst_nodes
|
||||
)
|
||||
|
||||
for ntype, nodes in original_row_ids.items():
|
||||
expected_nodes = expected_original_row_ids[ntype]
|
||||
assert torch.equal(nodes, expected_nodes)
|
||||
for etype, csc_format in compacted_csc_formats.items():
|
||||
indptr = csc_format.indptr
|
||||
indices = csc_format.indices
|
||||
expected_indptr = expected_csc_formats[etype].indptr
|
||||
expected_indices = expected_csc_formats[etype].indices
|
||||
assert torch.equal(indptr, expected_indptr)
|
||||
assert torch.equal(indices, expected_indices)
|
||||
|
||||
|
||||
def test_compact_csc_format_homo():
|
||||
seeds = torch.tensor([1, 3, 5, 2, 6])
|
||||
indptr = torch.tensor([0, 2, 4, 6, 7, 11])
|
||||
indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
|
||||
csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
|
||||
|
||||
expected_original_row_ids = torch.tensor(
|
||||
[1, 3, 5, 2, 6, 2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6]
|
||||
)
|
||||
expected_indptr = indptr
|
||||
expected_indices = torch.arange(0, len(indices)) + 5
|
||||
|
||||
original_row_ids, compacted_csc_formats = gb.compact_csc_format(
|
||||
csc_formats, seeds
|
||||
)
|
||||
|
||||
indptr = compacted_csc_formats.indptr
|
||||
indices = compacted_csc_formats.indices
|
||||
|
||||
assert torch.equal(indptr, expected_indptr)
|
||||
assert torch.equal(indices, expected_indices)
|
||||
assert torch.equal(original_row_ids, expected_original_row_ids)
|
||||
|
||||
|
||||
def test_compact_incorrect_indptr():
|
||||
seeds = torch.tensor([1, 3, 5, 2, 6, 7])
|
||||
indptr = torch.tensor([0, 2, 4, 6, 7, 11])
|
||||
indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4, 6])
|
||||
csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
|
||||
|
||||
# The number of seeds is not corresponding to indptr.
|
||||
with pytest.raises(AssertionError):
|
||||
gb.compact_csc_format(csc_formats, seeds)
|
||||
@@ -0,0 +1,286 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from functools import partial
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
|
||||
import dgl.graphbolt.internal as internal
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def test_read_torch_data():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
save_tensor = torch.tensor([[1, 2, 4], [2, 5, 3]])
|
||||
file_name = os.path.join(test_dir, "save_tensor.pt")
|
||||
torch.save(save_tensor, file_name)
|
||||
read_tensor = internal.utils._read_torch_data(file_name)
|
||||
assert torch.equal(save_tensor, read_tensor)
|
||||
save_tensor = read_tensor = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("in_memory", [True, False])
|
||||
def test_read_numpy_data(in_memory):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
save_numpy = np.array([[1, 2, 4], [2, 5, 3]])
|
||||
file_name = os.path.join(test_dir, "save_numpy.npy")
|
||||
np.save(file_name, save_numpy)
|
||||
read_tensor = internal.utils._read_numpy_data(file_name, in_memory)
|
||||
assert torch.equal(torch.from_numpy(save_numpy), read_tensor)
|
||||
save_numpy = read_tensor = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["torch", "numpy"])
|
||||
def test_read_data(fmt):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
data = np.array([[1, 2, 4], [2, 5, 3]])
|
||||
type_name = "pt" if fmt == "torch" else "npy"
|
||||
file_name = os.path.join(test_dir, f"save_data.{type_name}")
|
||||
if fmt == "numpy":
|
||||
np.save(file_name, data)
|
||||
elif fmt == "torch":
|
||||
torch.save(torch.from_numpy(data), file_name)
|
||||
read_tensor = internal.read_data(file_name, fmt)
|
||||
assert torch.equal(torch.from_numpy(data), read_tensor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_fmt, save_fmt, contiguous",
|
||||
[
|
||||
("torch", "torch", True),
|
||||
("torch", "torch", False),
|
||||
("torch", "numpy", True),
|
||||
("torch", "numpy", False),
|
||||
("numpy", "torch", True),
|
||||
("numpy", "torch", False),
|
||||
("numpy", "numpy", True),
|
||||
("numpy", "numpy", False),
|
||||
],
|
||||
)
|
||||
def test_save_data(data_fmt, save_fmt, contiguous):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
data = np.array([[1, 2, 4], [2, 5, 3]])
|
||||
if not contiguous:
|
||||
data = np.asfortranarray(data)
|
||||
tensor_data = torch.from_numpy(data)
|
||||
type_name = "pt" if save_fmt == "torch" else "npy"
|
||||
save_file_name = os.path.join(test_dir, f"save_data.{type_name}")
|
||||
# Step1. Save the data.
|
||||
if data_fmt == "torch":
|
||||
internal.save_data(tensor_data, save_file_name, save_fmt)
|
||||
elif data_fmt == "numpy":
|
||||
internal.save_data(data, save_file_name, save_fmt)
|
||||
|
||||
# Step2. Load the data.
|
||||
if save_fmt == "torch":
|
||||
loaded_data = torch.load(save_file_name, weights_only=False)
|
||||
assert loaded_data.is_contiguous()
|
||||
assert torch.equal(tensor_data, loaded_data)
|
||||
elif save_fmt == "numpy":
|
||||
loaded_data = np.load(save_file_name)
|
||||
# Checks if the loaded data is C-contiguous.
|
||||
assert loaded_data.flags["C_CONTIGUOUS"]
|
||||
assert np.array_equal(tensor_data.numpy(), loaded_data)
|
||||
|
||||
data = tensor_data = loaded_data = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["torch", "numpy"])
|
||||
def test_get_npy_dim(fmt):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
data = np.array([[1, 2, 4], [2, 5, 3]])
|
||||
type_name = "pt" if fmt == "torch" else "npy"
|
||||
file_name = os.path.join(test_dir, f"save_data.{type_name}")
|
||||
if fmt == "numpy":
|
||||
np.save(file_name, data)
|
||||
assert internal.get_npy_dim(file_name) == 2
|
||||
elif fmt == "torch":
|
||||
torch.save(torch.from_numpy(data), file_name)
|
||||
with pytest.raises(ValueError):
|
||||
internal.get_npy_dim(file_name)
|
||||
data = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_fmt", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize("save_fmt", ["numpy", "torch"])
|
||||
@pytest.mark.parametrize("is_feature", [True, False])
|
||||
def test_copy_or_convert_data(data_fmt, save_fmt, is_feature):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
data = np.arange(10)
|
||||
tensor_data = torch.from_numpy(data)
|
||||
in_type_name = "npy" if data_fmt == "numpy" else "pt"
|
||||
input_path = os.path.join(test_dir, f"data.{in_type_name}")
|
||||
out_type_name = "npy" if save_fmt == "numpy" else "pt"
|
||||
output_path = os.path.join(test_dir, f"out_data.{out_type_name}")
|
||||
if data_fmt == "numpy":
|
||||
np.save(input_path, data)
|
||||
else:
|
||||
torch.save(tensor_data, input_path)
|
||||
if save_fmt == "torch":
|
||||
with pytest.raises(AssertionError):
|
||||
internal.copy_or_convert_data(
|
||||
input_path,
|
||||
output_path,
|
||||
data_fmt,
|
||||
save_fmt,
|
||||
is_feature=is_feature,
|
||||
)
|
||||
else:
|
||||
internal.copy_or_convert_data(
|
||||
input_path,
|
||||
output_path,
|
||||
data_fmt,
|
||||
save_fmt,
|
||||
is_feature=is_feature,
|
||||
)
|
||||
if is_feature:
|
||||
data = data.reshape(-1, 1)
|
||||
tensor_data = tensor_data.reshape(-1, 1)
|
||||
if save_fmt == "numpy":
|
||||
out_data = np.load(output_path)
|
||||
assert (data == out_data).all()
|
||||
|
||||
data = None
|
||||
tensor_data = None
|
||||
out_data = None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("edge_fmt", ["csv", "numpy"])
|
||||
def test_read_edges(edge_fmt):
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
num_nodes = 40
|
||||
num_edges = 200
|
||||
nodes = np.repeat(np.arange(num_nodes), 5)
|
||||
neighbors = np.random.randint(0, num_nodes, size=(num_edges))
|
||||
edges = np.stack([nodes, neighbors], axis=1)
|
||||
os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
|
||||
if edge_fmt == "csv":
|
||||
# Wrtie into edges/edge.csv
|
||||
edges = pd.DataFrame(edges, columns=["src", "dst"])
|
||||
edge_path = os.path.join("edges", "edge.csv")
|
||||
edges.to_csv(
|
||||
os.path.join(test_dir, edge_path),
|
||||
index=False,
|
||||
header=False,
|
||||
)
|
||||
else:
|
||||
# Wrtie into edges/edge.npy
|
||||
edges = edges.T
|
||||
edge_path = os.path.join("edges", "edge.npy")
|
||||
np.save(os.path.join(test_dir, edge_path), edges)
|
||||
src, dst = internal.read_edges(test_dir, edge_fmt, edge_path)
|
||||
assert src.all() == nodes.all()
|
||||
assert dst.all() == neighbors.all()
|
||||
|
||||
|
||||
def test_read_edges_error():
|
||||
# 1. Unsupported file format.
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match="`numpy` or `csv` is expected when reading edges but got `fake-type`.",
|
||||
):
|
||||
internal.read_edges("test_dir", "fake-type", "edge_path")
|
||||
|
||||
# 2. Unexpected shape of numpy array
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
num_nodes = 40
|
||||
num_edges = 200
|
||||
nodes = np.repeat(np.arange(num_nodes), 5)
|
||||
neighbors = np.random.randint(0, num_nodes, size=(num_edges))
|
||||
edges = np.stack([nodes, neighbors, nodes], axis=1)
|
||||
os.makedirs(os.path.join(test_dir, "edges"), exist_ok=True)
|
||||
# Wrtie into edges/edge.npy
|
||||
edges = edges.T
|
||||
edge_path = os.path.join("edges", "edge.npy")
|
||||
np.save(os.path.join(test_dir, edge_path), edges)
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"The shape of edges should be (2, N), but got torch.Size([3, 200])."
|
||||
),
|
||||
):
|
||||
internal.read_edges(test_dir, "numpy", edge_path)
|
||||
|
||||
|
||||
def test_calculate_file_hash():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
test_file_path = os.path.join(test_dir, "test.txt")
|
||||
with open(test_file_path, "w") as file:
|
||||
file.write("test content")
|
||||
hash_value = internal.calculate_file_hash(
|
||||
test_file_path, hash_algo="md5"
|
||||
)
|
||||
expected_hash_value = "9473fdd0d880a43c21b7778d34872157"
|
||||
assert expected_hash_value == hash_value
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=re.escape(
|
||||
"Hash algorithm must be one of: ['md5', 'sha1', 'sha224', "
|
||||
+ "'sha256', 'sha384', 'sha512'], but got `fake`."
|
||||
),
|
||||
):
|
||||
hash_value = internal.calculate_file_hash(
|
||||
test_file_path, hash_algo="fake"
|
||||
)
|
||||
|
||||
|
||||
def test_calculate_dir_hash():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
test_file_path_1 = os.path.join(test_dir, "test_1.txt")
|
||||
test_file_path_2 = os.path.join(test_dir, "test_2.txt")
|
||||
with open(test_file_path_1, "w") as file:
|
||||
file.write("test content")
|
||||
with open(test_file_path_2, "w") as file:
|
||||
file.write("test contents of directory")
|
||||
hash_value = internal.calculate_dir_hash(test_dir, hash_algo="md5")
|
||||
expected_hash_value = [
|
||||
"56e708a2bdf92887d4a7f25cbc13c555",
|
||||
"9473fdd0d880a43c21b7778d34872157",
|
||||
]
|
||||
assert len(hash_value) == 2
|
||||
for val in hash_value.values():
|
||||
assert val in expected_hash_value
|
||||
|
||||
|
||||
def test_check_dataset_change():
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
# Generate directory and record its hash value.
|
||||
test_file_path_1 = os.path.join(test_dir, "test_1.txt")
|
||||
test_file_path_2 = os.path.join(test_dir, "test_2.txt")
|
||||
with open(test_file_path_1, "w") as file:
|
||||
file.write("test content")
|
||||
with open(test_file_path_2, "w") as file:
|
||||
file.write("test contents of directory")
|
||||
hash_value = internal.calculate_dir_hash(test_dir, hash_algo="md5")
|
||||
hash_value_file = "dataset_hash_value.txt"
|
||||
hash_value_file_paht = os.path.join(
|
||||
test_dir, "preprocessed", hash_value_file
|
||||
)
|
||||
os.makedirs(os.path.join(test_dir, "preprocessed"), exist_ok=True)
|
||||
with open(hash_value_file_paht, "w") as file:
|
||||
file.write(json.dumps(hash_value, indent=4))
|
||||
|
||||
# Modify the content of a file.
|
||||
with open(test_file_path_2, "w") as file:
|
||||
file.write("test contents of directory changed")
|
||||
|
||||
assert internal.check_dataset_change(test_dir, "preprocessed")
|
||||
|
||||
|
||||
def test_numpy_save_aligned():
|
||||
assert_equal = partial(torch.testing.assert_close, rtol=0, atol=0)
|
||||
a = torch.randn(1024, dtype=torch.float32) # 4096 bytes
|
||||
with tempfile.TemporaryDirectory() as test_dir:
|
||||
aligned_path = os.path.join(test_dir, "aligned.npy")
|
||||
gb.numpy_save_aligned(aligned_path, a.numpy())
|
||||
|
||||
nonaligned_path = os.path.join(test_dir, "nonaligned.npy")
|
||||
np.save(nonaligned_path, a.numpy())
|
||||
|
||||
assert_equal(np.load(aligned_path), np.load(nonaligned_path))
|
||||
# The size of the file should be 4K (aligned header) + 4K (tensor).
|
||||
assert os.path.getsize(aligned_path) == 4096 * 2
|
||||
@@ -0,0 +1,413 @@
|
||||
import os
|
||||
import re
|
||||
import unittest
|
||||
from collections.abc import Iterable, Mapping
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
from torch.torch_version import TorchVersion
|
||||
|
||||
from . import gb_test_utils
|
||||
|
||||
|
||||
def test_pytorch_cuda_allocator_conf():
|
||||
env = os.getenv("PYTORCH_CUDA_ALLOC_CONF")
|
||||
assert env is not None
|
||||
config_list = env.split(",")
|
||||
assert "expandable_segments:True" in config_list
|
||||
|
||||
|
||||
@unittest.skipIf(F._default_context_str != "gpu", "CopyTo needs GPU to test")
|
||||
@pytest.mark.parametrize("non_blocking", [False, True])
|
||||
def test_CopyTo(non_blocking):
|
||||
item_sampler = gb.ItemSampler(
|
||||
gb.ItemSet(torch.arange(20), names="seeds"), 4
|
||||
)
|
||||
if non_blocking:
|
||||
item_sampler = item_sampler.transform(lambda x: x.pin_memory())
|
||||
|
||||
# Invoke CopyTo via class constructor.
|
||||
dp = gb.CopyTo(item_sampler, "cuda")
|
||||
for data in dp:
|
||||
assert data.seeds.device.type == "cuda"
|
||||
|
||||
dp = gb.CopyTo(item_sampler, "cuda", non_blocking)
|
||||
for data in dp:
|
||||
assert data.seeds.device.type == "cuda"
|
||||
|
||||
# Invoke CopyTo via functional form.
|
||||
dp = item_sampler.copy_to("cuda", non_blocking)
|
||||
for data in dp:
|
||||
assert data.seeds.device.type == "cuda"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"task",
|
||||
[
|
||||
"node_classification",
|
||||
"node_inference",
|
||||
"link_prediction",
|
||||
"edge_classification",
|
||||
],
|
||||
)
|
||||
@unittest.skipIf(F._default_context_str == "cpu", "CopyTo needs GPU to test")
|
||||
def test_CopyToWithMiniBatches(task):
|
||||
N = 16
|
||||
B = 2
|
||||
if task == "node_classification":
|
||||
itemset = gb.ItemSet(
|
||||
(torch.arange(N), torch.arange(N)), names=("seeds", "labels")
|
||||
)
|
||||
elif task == "node_inference":
|
||||
itemset = gb.ItemSet(torch.arange(N), names="seeds")
|
||||
elif task == "link_prediction":
|
||||
itemset = gb.ItemSet(
|
||||
(
|
||||
torch.arange(2 * N).reshape(-1, 2),
|
||||
torch.arange(N),
|
||||
),
|
||||
names=("seeds", "labels"),
|
||||
)
|
||||
elif task == "edge_classification":
|
||||
itemset = gb.ItemSet(
|
||||
(torch.arange(2 * N).reshape(-1, 2), torch.arange(N)),
|
||||
names=("seeds", "labels"),
|
||||
)
|
||||
graph = gb_test_utils.rand_csc_graph(100, 0.15, bidirection_edge=True)
|
||||
|
||||
features = {}
|
||||
keys = [("node", None, "a"), ("node", None, "b")]
|
||||
features[keys[0]] = gb.TorchBasedFeature(torch.randn(200, 4))
|
||||
features[keys[1]] = gb.TorchBasedFeature(torch.randn(200, 4))
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
datapipe = gb.ItemSampler(itemset, batch_size=B)
|
||||
datapipe = gb.NeighborSampler(
|
||||
datapipe,
|
||||
graph,
|
||||
fanouts=[torch.LongTensor([2]) for _ in range(2)],
|
||||
)
|
||||
if task != "node_inference":
|
||||
datapipe = gb.FeatureFetcher(
|
||||
datapipe,
|
||||
feature_store,
|
||||
["a"],
|
||||
)
|
||||
|
||||
copied_attrs = [
|
||||
"labels",
|
||||
"compacted_seeds",
|
||||
"sampled_subgraphs",
|
||||
"indexes",
|
||||
"node_features",
|
||||
"edge_features",
|
||||
"blocks",
|
||||
"seeds",
|
||||
"input_nodes",
|
||||
]
|
||||
|
||||
def test_data_device(datapipe):
|
||||
for data in datapipe:
|
||||
for attr in dir(data):
|
||||
var = getattr(data, attr)
|
||||
if isinstance(var, Mapping):
|
||||
var = var[next(iter(var))]
|
||||
elif isinstance(var, Iterable):
|
||||
var = next(iter(var))
|
||||
if (
|
||||
not callable(var)
|
||||
and not attr.startswith("__")
|
||||
and hasattr(var, "device")
|
||||
and var is not None
|
||||
):
|
||||
if attr in copied_attrs:
|
||||
assert var.device.type == "cuda", attr
|
||||
else:
|
||||
assert var.device.type == "cpu", attr
|
||||
|
||||
# Invoke CopyTo via class constructor.
|
||||
test_data_device(gb.CopyTo(datapipe, "cuda"))
|
||||
|
||||
# Invoke CopyTo via functional form.
|
||||
test_data_device(datapipe.copy_to("cuda"))
|
||||
|
||||
|
||||
def test_etype_tuple_to_str():
|
||||
"""Convert etype from tuple to string."""
|
||||
# Test for expected input.
|
||||
c_etype = ("user", "like", "item")
|
||||
c_etype_str = gb.etype_tuple_to_str(c_etype)
|
||||
assert c_etype_str == "user:like:item"
|
||||
|
||||
# Test for unexpected input: not a tuple.
|
||||
c_etype = "user:like:item"
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Passed-in canonical etype should be in format of (str, str, str). "
|
||||
"But got user:like:item."
|
||||
),
|
||||
):
|
||||
_ = gb.etype_tuple_to_str(c_etype)
|
||||
|
||||
# Test for unexpected input: tuple with wrong length.
|
||||
c_etype = ("user", "like")
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Passed-in canonical etype should be in format of (str, str, str). "
|
||||
"But got ('user', 'like')."
|
||||
),
|
||||
):
|
||||
_ = gb.etype_tuple_to_str(c_etype)
|
||||
|
||||
|
||||
def test_etype_str_to_tuple():
|
||||
"""Convert etype from string to tuple."""
|
||||
# Test for expected input.
|
||||
c_etype_str = "user:like:item"
|
||||
c_etype = gb.etype_str_to_tuple(c_etype_str)
|
||||
assert c_etype == ("user", "like", "item")
|
||||
|
||||
# Test for unexpected input: string with wrong format.
|
||||
c_etype_str = "user:like"
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Passed-in canonical etype should be in format of 'str:str:str'. "
|
||||
"But got user:like."
|
||||
),
|
||||
):
|
||||
_ = gb.etype_str_to_tuple(c_etype_str)
|
||||
|
||||
|
||||
def test_seed_type_str_to_ntypes():
|
||||
"""Convert etype from string to tuple."""
|
||||
# Test for node pairs.
|
||||
seed_type_str = "user:like:item"
|
||||
seed_size = 2
|
||||
node_type = gb.seed_type_str_to_ntypes(seed_type_str, seed_size)
|
||||
assert node_type == ["user", "item"]
|
||||
|
||||
# Test for node pairs.
|
||||
seed_type_str = "user:item:user"
|
||||
seed_size = 3
|
||||
node_type = gb.seed_type_str_to_ntypes(seed_type_str, seed_size)
|
||||
assert node_type == ["user", "item", "user"]
|
||||
|
||||
# Test for unexpected input: list.
|
||||
seed_type_str = ["user", "item"]
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Passed-in seed type should be string, but got <class 'list'>"
|
||||
),
|
||||
):
|
||||
_ = gb.seed_type_str_to_ntypes(seed_type_str, 2)
|
||||
|
||||
|
||||
def test_isin():
|
||||
elements = torch.tensor([2, 3, 5, 5, 20, 13, 11], device=F.ctx())
|
||||
test_elements = torch.tensor([2, 5], device=F.ctx())
|
||||
res = gb.isin(elements, test_elements)
|
||||
expected = torch.tensor(
|
||||
[True, False, True, True, False, False, False], device=F.ctx()
|
||||
)
|
||||
assert torch.equal(res, expected)
|
||||
|
||||
|
||||
def test_isin_big_data():
|
||||
elements = torch.randint(0, 10000, (10000000,), device=F.ctx())
|
||||
test_elements = torch.randint(0, 10000, (500000,), device=F.ctx())
|
||||
res = gb.isin(elements, test_elements)
|
||||
expected = torch.isin(elements, test_elements)
|
||||
assert torch.equal(res, expected)
|
||||
|
||||
|
||||
def test_isin_non_1D_dim():
|
||||
elements = torch.tensor([[2, 3], [5, 5], [20, 13]], device=F.ctx())
|
||||
test_elements = torch.tensor([2, 5], device=F.ctx())
|
||||
with pytest.raises(Exception):
|
||||
gb.isin(elements, test_elements)
|
||||
elements = torch.tensor([2, 3, 5, 5, 20, 13], device=F.ctx())
|
||||
test_elements = torch.tensor([[2, 5]], device=F.ctx())
|
||||
with pytest.raises(Exception):
|
||||
gb.isin(elements, test_elements)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("pinned", [False, True])
|
||||
def test_index_select(dtype, idtype, pinned):
|
||||
if F._default_context_str != "gpu" and pinned:
|
||||
pytest.skip("Pinned tests are available only on GPU.")
|
||||
tensor = torch.tensor([[2, 3], [5, 5], [20, 13]], dtype=dtype)
|
||||
tensor = tensor.pin_memory() if pinned else tensor.to(F.ctx())
|
||||
index = torch.tensor([0, 2], dtype=idtype, device=F.ctx())
|
||||
gb_result = gb.index_select(tensor, index)
|
||||
torch_result = tensor.to(F.ctx())[index.long()]
|
||||
assert torch.equal(torch_result, gb_result)
|
||||
if pinned:
|
||||
gb_result = gb.index_select(tensor.cpu(), index.cpu().pin_memory())
|
||||
assert torch.equal(torch_result.cpu(), gb_result)
|
||||
assert gb_result.is_pinned()
|
||||
|
||||
# Test the internal async API
|
||||
future = torch.ops.graphbolt.index_select_async(tensor.cpu(), index.cpu())
|
||||
assert torch.equal(torch_result.cpu(), future.wait())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dtype",
|
||||
[
|
||||
torch.bool,
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.int16,
|
||||
torch.int32,
|
||||
torch.int64,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float32,
|
||||
torch.float64,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("idtype", [torch.int32, torch.int64])
|
||||
def test_scatter_async(dtype, idtype):
|
||||
input = torch.tensor([[2, 3], [5, 5], [20, 13]], dtype=dtype)
|
||||
index = torch.ones([1], dtype=idtype)
|
||||
res = torch.ops.graphbolt.scatter_async(input, index, input[2:3])
|
||||
assert torch.equal(
|
||||
torch.tensor([[2, 3], [20, 13], [20, 13]], dtype=dtype), res.wait()
|
||||
)
|
||||
|
||||
|
||||
def torch_expand_indptr(indptr, dtype, nodes=None):
|
||||
if nodes is None:
|
||||
nodes = torch.arange(len(indptr) - 1, dtype=dtype, device=indptr.device)
|
||||
return nodes.to(dtype).repeat_interleave(indptr.diff())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nodes", [None, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
def test_expand_indptr(nodes, dtype):
|
||||
if nodes:
|
||||
nodes = torch.tensor([1, 7, 3, 4, 5, 8], dtype=dtype, device=F.ctx())
|
||||
indptr = torch.tensor([0, 2, 2, 7, 10, 12, 20], device=F.ctx())
|
||||
torch_result = torch_expand_indptr(indptr, dtype, nodes)
|
||||
gb_result = gb.expand_indptr(indptr, dtype, nodes)
|
||||
assert torch.equal(torch_result, gb_result)
|
||||
gb_result = gb.expand_indptr(indptr, dtype, nodes, indptr[-1].item())
|
||||
assert torch.equal(torch_result, gb_result)
|
||||
|
||||
if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
|
||||
import torch._dynamo as dynamo
|
||||
from torch.testing._internal.optests import opcheck
|
||||
|
||||
# Tests torch.compile compatibility
|
||||
for output_size in [None, indptr[-1].item()]:
|
||||
kwargs = {"node_ids": nodes, "output_size": output_size}
|
||||
opcheck(
|
||||
torch.ops.graphbolt.expand_indptr,
|
||||
(indptr, dtype),
|
||||
kwargs,
|
||||
test_utils=[
|
||||
"test_schema",
|
||||
"test_autograd_registration",
|
||||
"test_faketensor",
|
||||
"test_aot_dispatch_dynamic",
|
||||
],
|
||||
raise_exception=True,
|
||||
)
|
||||
|
||||
explanation = dynamo.explain(gb.expand_indptr)(
|
||||
indptr, dtype, nodes, output_size
|
||||
)
|
||||
expected_breaks = -1 if output_size is None else 0
|
||||
assert explanation.graph_break_count == expected_breaks
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str != "gpu", "Only GPU implementation is available."
|
||||
)
|
||||
@pytest.mark.parametrize("offset", [None, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
def test_indptr_edge_ids(offset, dtype):
|
||||
indptr = torch.tensor([0, 2, 2, 7, 10, 12], device=F.ctx())
|
||||
if offset:
|
||||
offset = indptr[:-1]
|
||||
ref_result = torch.arange(
|
||||
0, indptr[-1].item(), dtype=dtype, device=F.ctx()
|
||||
)
|
||||
else:
|
||||
ref_result = torch.tensor(
|
||||
[0, 1, 0, 1, 2, 3, 4, 0, 1, 2, 0, 1], dtype=dtype, device=F.ctx()
|
||||
)
|
||||
gb_result = gb.indptr_edge_ids(indptr, dtype, offset)
|
||||
assert torch.equal(ref_result, gb_result)
|
||||
gb_result = gb.indptr_edge_ids(indptr, dtype, offset, indptr[-1].item())
|
||||
assert torch.equal(ref_result, gb_result)
|
||||
|
||||
if TorchVersion(torch.__version__) >= TorchVersion("2.2.0a0"):
|
||||
import torch._dynamo as dynamo
|
||||
from torch.testing._internal.optests import opcheck
|
||||
|
||||
# Tests torch.compile compatibility
|
||||
for output_size in [None, indptr[-1].item()]:
|
||||
kwargs = {"offset": offset, "output_size": output_size}
|
||||
opcheck(
|
||||
torch.ops.graphbolt.indptr_edge_ids,
|
||||
(indptr, dtype),
|
||||
kwargs,
|
||||
test_utils=[
|
||||
"test_schema",
|
||||
"test_autograd_registration",
|
||||
"test_faketensor",
|
||||
"test_aot_dispatch_dynamic",
|
||||
],
|
||||
raise_exception=True,
|
||||
)
|
||||
|
||||
explanation = dynamo.explain(gb.indptr_edge_ids)(
|
||||
indptr, dtype, offset, output_size
|
||||
)
|
||||
expected_breaks = -1 if output_size is None else 0
|
||||
assert explanation.graph_break_count == expected_breaks
|
||||
|
||||
|
||||
def test_csc_format_base_representation():
|
||||
csc_format_base = gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4]),
|
||||
indices=torch.tensor([4, 5, 6, 7]),
|
||||
)
|
||||
expected_result = str(
|
||||
"""CSCFormatBase(indptr=tensor([0, 2, 4]),
|
||||
indices=tensor([4, 5, 6, 7]),
|
||||
)"""
|
||||
)
|
||||
assert str(csc_format_base) == expected_result, print(csc_format_base)
|
||||
|
||||
|
||||
def test_csc_format_base_incorrect_indptr():
|
||||
indptr = torch.tensor([0, 2, 4, 6, 7, 11])
|
||||
indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
|
||||
with pytest.raises(AssertionError):
|
||||
# The value of last element in indptr is not corresponding to indices.
|
||||
csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
|
||||
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
import unittest
|
||||
from sys import platform
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import dgl.graphbolt
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as thd
|
||||
|
||||
from dgl.graphbolt.datapipes import find_dps, traverse_dps
|
||||
|
||||
from . import gb_test_utils
|
||||
|
||||
|
||||
@pytest.mark.parametrize("overlap_feature_fetch", [False, True])
|
||||
def test_DataLoader(overlap_feature_fetch):
|
||||
N = 40
|
||||
B = 4
|
||||
itemset = dgl.graphbolt.ItemSet(torch.arange(N), names="seeds")
|
||||
graph = gb_test_utils.rand_csc_graph(200, 0.15, bidirection_edge=True)
|
||||
features = {}
|
||||
keys = [("node", None, "a"), ("node", None, "b"), ("edge", None, "c")]
|
||||
features[keys[0]] = dgl.graphbolt.TorchBasedFeature(torch.randn(200, 4))
|
||||
features[keys[1]] = dgl.graphbolt.TorchBasedFeature(torch.randn(200, 4))
|
||||
M = graph.total_num_edges
|
||||
features[keys[2]] = dgl.graphbolt.TorchBasedFeature(torch.randn(M, 1))
|
||||
feature_store = dgl.graphbolt.BasicFeatureStore(features)
|
||||
|
||||
item_sampler = dgl.graphbolt.ItemSampler(itemset, batch_size=B)
|
||||
subgraph_sampler = dgl.graphbolt.NeighborSampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
fanouts=[torch.LongTensor([2]) for _ in range(2)],
|
||||
)
|
||||
feature_fetcher = dgl.graphbolt.FeatureFetcher(
|
||||
subgraph_sampler,
|
||||
feature_store,
|
||||
["a", "b"],
|
||||
["c"],
|
||||
overlap_fetch=overlap_feature_fetch,
|
||||
)
|
||||
device_transferrer = dgl.graphbolt.CopyTo(feature_fetcher, F.ctx())
|
||||
|
||||
dataloader = dgl.graphbolt.DataLoader(
|
||||
device_transferrer,
|
||||
num_workers=4,
|
||||
)
|
||||
for i, minibatch in enumerate(dataloader):
|
||||
assert "a" in minibatch.node_features
|
||||
assert "b" in minibatch.node_features
|
||||
for layer_id in range(minibatch.num_layers()):
|
||||
assert "c" in minibatch.edge_features[layer_id]
|
||||
assert i + 1 == N // B
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str != "gpu",
|
||||
reason="This test requires the GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"sampler_name", ["NeighborSampler", "LayerNeighborSampler"]
|
||||
)
|
||||
@pytest.mark.parametrize("enable_feature_fetch", [True, False])
|
||||
@pytest.mark.parametrize("overlap_feature_fetch", [True, False])
|
||||
@pytest.mark.parametrize("overlap_graph_fetch", [True, False])
|
||||
@pytest.mark.parametrize("cooperative", [True, False])
|
||||
@pytest.mark.parametrize("asynchronous", [True, False])
|
||||
@pytest.mark.parametrize("num_gpu_cached_edges", [0, 1024])
|
||||
@pytest.mark.parametrize("gpu_cache_threshold", [1, 3])
|
||||
def test_gpu_sampling_DataLoader(
|
||||
sampler_name,
|
||||
enable_feature_fetch,
|
||||
overlap_feature_fetch,
|
||||
overlap_graph_fetch,
|
||||
cooperative,
|
||||
asynchronous,
|
||||
num_gpu_cached_edges,
|
||||
gpu_cache_threshold,
|
||||
):
|
||||
if cooperative and not thd.is_initialized():
|
||||
# On Windows, the init method can only be file.
|
||||
init_method = (
|
||||
f"file:///{os.path.join(os.getcwd(), 'dis_tempfile')}"
|
||||
if platform == "win32"
|
||||
else "tcp://127.0.0.1:12345"
|
||||
)
|
||||
thd.init_process_group(
|
||||
init_method=init_method,
|
||||
world_size=1,
|
||||
rank=0,
|
||||
)
|
||||
N = 40
|
||||
B = 4
|
||||
num_layers = 2
|
||||
itemset = dgl.graphbolt.ItemSet(torch.arange(N), names="seeds")
|
||||
graph = gb_test_utils.rand_csc_graph(200, 0.15, bidirection_edge=True)
|
||||
graph = graph.pin_memory_() if overlap_graph_fetch else graph.to(F.ctx())
|
||||
features = {}
|
||||
keys = [
|
||||
("node", None, "a"),
|
||||
("node", None, "b"),
|
||||
("node", None, "c"),
|
||||
("edge", None, "d"),
|
||||
]
|
||||
features[keys[0]] = dgl.graphbolt.TorchBasedFeature(
|
||||
torch.randn(200, 4, pin_memory=True)
|
||||
)
|
||||
features[keys[1]] = dgl.graphbolt.TorchBasedFeature(
|
||||
torch.randn(200, 4, pin_memory=True)
|
||||
)
|
||||
features[keys[2]] = dgl.graphbolt.TorchBasedFeature(
|
||||
torch.randn(200, 4, device=F.ctx())
|
||||
)
|
||||
features[keys[3]] = dgl.graphbolt.TorchBasedFeature(
|
||||
torch.randn(graph.total_num_edges, 1, device=F.ctx())
|
||||
)
|
||||
feature_store = dgl.graphbolt.BasicFeatureStore(features)
|
||||
|
||||
dataloaders = []
|
||||
for i in range(2):
|
||||
datapipe = dgl.graphbolt.ItemSampler(itemset, batch_size=B)
|
||||
datapipe = datapipe.copy_to(F.ctx())
|
||||
kwargs = {
|
||||
"overlap_fetch": overlap_graph_fetch,
|
||||
"num_gpu_cached_edges": num_gpu_cached_edges,
|
||||
"gpu_cache_threshold": gpu_cache_threshold,
|
||||
"cooperative": cooperative,
|
||||
"asynchronous": asynchronous,
|
||||
}
|
||||
if i != 0:
|
||||
kwargs = {}
|
||||
datapipe = getattr(dgl.graphbolt, sampler_name)(
|
||||
datapipe,
|
||||
graph,
|
||||
fanouts=[torch.LongTensor([2]) for _ in range(num_layers)],
|
||||
**kwargs,
|
||||
)
|
||||
if enable_feature_fetch:
|
||||
datapipe = dgl.graphbolt.FeatureFetcher(
|
||||
datapipe,
|
||||
feature_store,
|
||||
["a", "b", "c"],
|
||||
["d"],
|
||||
overlap_fetch=overlap_feature_fetch and i == 0,
|
||||
cooperative=asynchronous and cooperative and i == 0,
|
||||
)
|
||||
dataloaders.append(dgl.graphbolt.DataLoader(datapipe))
|
||||
dataloader, dataloader2 = dataloaders
|
||||
|
||||
bufferer_cnt = int(enable_feature_fetch and overlap_feature_fetch)
|
||||
if overlap_graph_fetch:
|
||||
bufferer_cnt += num_layers
|
||||
if num_gpu_cached_edges > 0:
|
||||
bufferer_cnt += 2 * num_layers
|
||||
if asynchronous:
|
||||
bufferer_cnt += 2 * num_layers + 1 # _preprocess stage has 1.
|
||||
if cooperative:
|
||||
bufferer_cnt += 3 * num_layers
|
||||
if enable_feature_fetch:
|
||||
bufferer_cnt += 1 # feature fetch has 1.
|
||||
if cooperative:
|
||||
# _preprocess stage.
|
||||
bufferer_cnt += 4
|
||||
datapipe_graph = traverse_dps(dataloader)
|
||||
bufferers = find_dps(
|
||||
datapipe_graph,
|
||||
dgl.graphbolt.Bufferer,
|
||||
)
|
||||
assert len(bufferers) == bufferer_cnt
|
||||
# Fixes the randomness of LayerNeighborSampler
|
||||
torch.manual_seed(1)
|
||||
minibatches = list(dataloader)
|
||||
assert len(minibatches) == N // B
|
||||
|
||||
for i, _ in enumerate(dataloader):
|
||||
if i >= 1:
|
||||
break
|
||||
|
||||
torch.manual_seed(1)
|
||||
|
||||
for minibatch, minibatch2 in zip(minibatches, dataloader2):
|
||||
if enable_feature_fetch:
|
||||
assert "a" in minibatch.node_features
|
||||
assert "b" in minibatch.node_features
|
||||
assert "c" in minibatch.node_features
|
||||
if sampler_name == "LayerNeighborSampler":
|
||||
assert torch.equal(
|
||||
minibatch.node_features["a"], minibatch2.node_features["a"]
|
||||
)
|
||||
for layer_id in range(minibatch.num_layers()):
|
||||
assert "d" in minibatch.edge_features[layer_id]
|
||||
edge_feature = minibatch.edge_features[layer_id]["d"]
|
||||
edge_feature_ref = minibatch2.edge_features[layer_id]["d"]
|
||||
if sampler_name == "LayerNeighborSampler":
|
||||
assert torch.equal(edge_feature, edge_feature_ref)
|
||||
assert len(list(dataloader)) == N // B
|
||||
|
||||
if asynchronous and cooperative:
|
||||
for minibatch in minibatches:
|
||||
x = torch.ones((minibatch.node_ids().size(0), 1), device=F.ctx())
|
||||
for subgraph in minibatch.sampled_subgraphs:
|
||||
x = gb.CooperativeConvFunction.apply(subgraph, x)
|
||||
x, edge_index, size = subgraph.to_pyg(x)
|
||||
x = x[0]
|
||||
one = torch.ones(
|
||||
edge_index.shape[1], dtype=x.dtype, device=x.device
|
||||
)
|
||||
coo = torch.sparse_coo_tensor(
|
||||
edge_index.flipud(), one, size=(size[1], size[0])
|
||||
)
|
||||
x = torch.sparse.mm(coo, x)
|
||||
assert x.shape[0] == minibatch.seeds.shape[0]
|
||||
assert x.shape[1] == 1
|
||||
|
||||
if thd.is_initialized():
|
||||
thd.destroy_process_group()
|
||||
@@ -0,0 +1,15 @@
|
||||
import pytest
|
||||
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def test_Dataset():
|
||||
dataset = gb.Dataset()
|
||||
with pytest.raises(NotImplementedError):
|
||||
_ = dataset.tasks
|
||||
with pytest.raises(NotImplementedError):
|
||||
_ = dataset.graph
|
||||
with pytest.raises(NotImplementedError):
|
||||
_ = dataset.feature
|
||||
with pytest.raises(NotImplementedError):
|
||||
_ = dataset.dataset_name
|
||||
@@ -0,0 +1,258 @@
|
||||
import random
|
||||
from functools import partial
|
||||
|
||||
import dgl.graphbolt as gb
|
||||
import torch
|
||||
from torch.utils.data.datapipes.iter import Mapper
|
||||
|
||||
from . import gb_test_utils
|
||||
|
||||
|
||||
def test_FeatureFetcher_invoke():
|
||||
# Prepare graph and required datapipes.
|
||||
graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
|
||||
a = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
|
||||
)
|
||||
b = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
|
||||
)
|
||||
|
||||
features = {}
|
||||
keys = [("node", None, "a"), ("edge", None, "b")]
|
||||
features[keys[0]] = gb.TorchBasedFeature(a)
|
||||
features[keys[1]] = gb.TorchBasedFeature(b)
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
itemset = gb.ItemSet(torch.arange(10), names="seeds")
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2)
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
|
||||
# Invoke FeatureFetcher via class constructor.
|
||||
datapipe = gb.NeighborSampler(item_sampler, graph, fanouts)
|
||||
|
||||
datapipe = gb.FeatureFetcher(datapipe, feature_store, ["a"], ["b"])
|
||||
assert len(list(datapipe)) == 5
|
||||
|
||||
# Invoke FeatureFetcher via functional form.
|
||||
datapipe = item_sampler.sample_neighbor(graph, fanouts).fetch_feature(
|
||||
feature_store, ["a"], ["b"]
|
||||
)
|
||||
assert len(list(datapipe)) == 5
|
||||
|
||||
|
||||
def test_FeatureFetcher_homo():
|
||||
graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
|
||||
a = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
|
||||
)
|
||||
b = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
|
||||
)
|
||||
|
||||
features = {}
|
||||
keys = [("node", None, "a"), ("edge", None, "b")]
|
||||
features[keys[0]] = gb.TorchBasedFeature(a)
|
||||
features[keys[1]] = gb.TorchBasedFeature(b)
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
itemset = gb.ItemSet(torch.arange(10), names="seeds")
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2)
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
sampler_dp = gb.NeighborSampler(item_sampler, graph, fanouts)
|
||||
fetcher_dp = gb.FeatureFetcher(sampler_dp, feature_store, ["a"], ["b"])
|
||||
|
||||
assert len(list(fetcher_dp)) == 5
|
||||
|
||||
|
||||
def _func(fn, minibatch):
|
||||
return fn(minibatch)
|
||||
|
||||
|
||||
def test_FeatureFetcher_with_edges_homo():
|
||||
graph = gb_test_utils.rand_csc_graph(20, 0.15, bidirection_edge=True)
|
||||
a = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_nodes)]
|
||||
)
|
||||
b = torch.tensor(
|
||||
[[random.randint(0, 10)] for _ in range(graph.total_num_edges)]
|
||||
)
|
||||
|
||||
def add_node_and_edge_ids(minibatch):
|
||||
seeds = minibatch.seeds
|
||||
subgraphs = []
|
||||
for _ in range(3):
|
||||
sampled_csc = gb.CSCFormatBase(
|
||||
indptr=torch.arange(11),
|
||||
indices=torch.arange(10),
|
||||
)
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=sampled_csc,
|
||||
original_column_node_ids=torch.arange(10),
|
||||
original_row_node_ids=torch.arange(10),
|
||||
original_edge_ids=torch.randint(
|
||||
0, graph.total_num_edges, (10,)
|
||||
),
|
||||
)
|
||||
)
|
||||
data = gb.MiniBatch(input_nodes=seeds, sampled_subgraphs=subgraphs)
|
||||
return data
|
||||
|
||||
features = {}
|
||||
keys = [("node", None, "a"), ("edge", None, "b")]
|
||||
features[keys[0]] = gb.TorchBasedFeature(a)
|
||||
features[keys[1]] = gb.TorchBasedFeature(b)
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
itemset = gb.ItemSet(torch.arange(10), names="seeds")
|
||||
item_sampler_dp = gb.ItemSampler(itemset, batch_size=2)
|
||||
fn = partial(_func, add_node_and_edge_ids)
|
||||
converter_dp = Mapper(item_sampler_dp, fn)
|
||||
fetcher_dp = gb.FeatureFetcher(converter_dp, feature_store, ["a"], ["b"])
|
||||
|
||||
assert len(list(fetcher_dp)) == 5
|
||||
for data in fetcher_dp:
|
||||
assert data.node_features["a"].size(0) == 2
|
||||
assert len(data.edge_features) == 3
|
||||
for edge_feature in data.edge_features:
|
||||
assert edge_feature["b"].size(0) == 10
|
||||
|
||||
|
||||
def get_hetero_graph():
|
||||
# COO graph:
|
||||
# [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
|
||||
# [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
|
||||
# [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
|
||||
# num_nodes = 5, num_n1 = 2, num_n2 = 3
|
||||
ntypes = {"n1": 0, "n2": 1}
|
||||
etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
|
||||
indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
|
||||
indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
|
||||
type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
|
||||
node_type_offset = torch.LongTensor([0, 2, 5])
|
||||
return gb.fused_csc_sampling_graph(
|
||||
indptr,
|
||||
indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
)
|
||||
|
||||
|
||||
def test_FeatureFetcher_hetero():
|
||||
graph = get_hetero_graph()
|
||||
a = torch.tensor([[random.randint(0, 10)] for _ in range(2)])
|
||||
b = torch.tensor([[random.randint(0, 10)] for _ in range(3)])
|
||||
|
||||
features = {}
|
||||
keys = [("node", "n1", "a"), ("node", "n2", "a")]
|
||||
features[keys[0]] = gb.TorchBasedFeature(a)
|
||||
features[keys[1]] = gb.TorchBasedFeature(b)
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
itemset = gb.HeteroItemSet(
|
||||
{
|
||||
"n1": gb.ItemSet(torch.LongTensor([0, 1]), names="seeds"),
|
||||
"n2": gb.ItemSet(torch.LongTensor([0, 1, 2]), names="seeds"),
|
||||
}
|
||||
)
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2)
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
sampler_dp = gb.NeighborSampler(item_sampler, graph, fanouts)
|
||||
# "n3" is not in the sampled input nodes.
|
||||
node_feature_keys = {"n1": ["a"], "n2": ["a"], "n3": ["a"]}
|
||||
fetcher_dp = gb.FeatureFetcher(
|
||||
sampler_dp, feature_store, node_feature_keys=node_feature_keys
|
||||
)
|
||||
assert len(list(fetcher_dp)) == 3
|
||||
|
||||
# Do not fetch feature for "n1".
|
||||
node_feature_keys = {"n2": ["a"]}
|
||||
fetcher_dp = gb.FeatureFetcher(
|
||||
sampler_dp, feature_store, node_feature_keys=node_feature_keys
|
||||
)
|
||||
for mini_batch in fetcher_dp:
|
||||
assert ("n1", "a") not in mini_batch.node_features
|
||||
|
||||
|
||||
def test_FeatureFetcher_with_edges_hetero():
|
||||
a = torch.tensor([[random.randint(0, 10)] for _ in range(20)])
|
||||
b = torch.tensor([[random.randint(0, 10)] for _ in range(50)])
|
||||
|
||||
def add_node_and_edge_ids(minibatch):
|
||||
seeds = minibatch.seeds
|
||||
subgraphs = []
|
||||
original_edge_ids = {
|
||||
"n1:e1:n2": torch.randint(0, 50, (10,)),
|
||||
"n2:e2:n1": torch.randint(0, 50, (10,)),
|
||||
}
|
||||
original_column_node_ids = {
|
||||
"n1": torch.randint(0, 20, (10,)),
|
||||
"n2": torch.randint(0, 20, (10,)),
|
||||
}
|
||||
original_row_node_ids = {
|
||||
"n1": torch.randint(0, 20, (10,)),
|
||||
"n2": torch.randint(0, 20, (10,)),
|
||||
}
|
||||
for _ in range(3):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc={
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.arange(11),
|
||||
indices=torch.arange(10),
|
||||
),
|
||||
"n2:e2:n1": gb.CSCFormatBase(
|
||||
indptr=torch.arange(11),
|
||||
indices=torch.arange(10),
|
||||
),
|
||||
},
|
||||
original_column_node_ids=original_column_node_ids,
|
||||
original_row_node_ids=original_row_node_ids,
|
||||
original_edge_ids=original_edge_ids,
|
||||
)
|
||||
)
|
||||
data = gb.MiniBatch(input_nodes=seeds, sampled_subgraphs=subgraphs)
|
||||
return data
|
||||
|
||||
features = {}
|
||||
keys = [
|
||||
("node", "n1", "a"),
|
||||
("edge", "n1:e1:n2", "a"),
|
||||
("edge", "n2:e2:n1", "a"),
|
||||
]
|
||||
features[keys[0]] = gb.TorchBasedFeature(a)
|
||||
features[keys[1]] = gb.TorchBasedFeature(b)
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
|
||||
itemset = gb.HeteroItemSet(
|
||||
{
|
||||
"n1": gb.ItemSet(torch.randint(0, 20, (10,)), names="seeds"),
|
||||
}
|
||||
)
|
||||
item_sampler_dp = gb.ItemSampler(itemset, batch_size=2)
|
||||
fn = partial(_func, add_node_and_edge_ids)
|
||||
converter_dp = Mapper(item_sampler_dp, fn)
|
||||
# "n3:e3:n3" is not in the sampled edges.
|
||||
# Do not fetch feature for "n2:e2:n1".
|
||||
node_feature_keys = {"n1": ["a"]}
|
||||
edge_feature_keys = {"n1:e1:n2": ["a"], "n3:e3:n3": ["a"]}
|
||||
fetcher_dp = gb.FeatureFetcher(
|
||||
converter_dp,
|
||||
feature_store,
|
||||
node_feature_keys=node_feature_keys,
|
||||
edge_feature_keys=edge_feature_keys,
|
||||
)
|
||||
|
||||
assert len(list(fetcher_dp)) == 5
|
||||
for data in fetcher_dp:
|
||||
assert data.node_features[("n1", "a")].size(0) == 2
|
||||
assert len(data.edge_features) == 3
|
||||
for edge_feature in data.edge_features:
|
||||
assert edge_feature[("n1:e1:n2", "a")].size(0) == 10
|
||||
assert ("n2:e2:n1", "a") not in edge_feature
|
||||
@@ -0,0 +1,60 @@
|
||||
import backend as F
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def test_find_reverse_edges_homo():
|
||||
edges = torch.tensor([[1, 3, 5], [2, 4, 5]]).T
|
||||
edges = gb.add_reverse_edges(edges)
|
||||
expected_edges = torch.tensor([[1, 3, 5, 2, 4, 5], [2, 4, 5, 1, 3, 5]]).T
|
||||
assert torch.equal(edges, expected_edges)
|
||||
assert torch.equal(edges[1], expected_edges[1])
|
||||
|
||||
|
||||
def test_find_reverse_edges_hetero():
|
||||
edges = {
|
||||
"A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
|
||||
"B:rr:A": torch.tensor([[3], [3]]).T,
|
||||
}
|
||||
edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A"})
|
||||
expected_edges = {
|
||||
"A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
|
||||
"B:rr:A": torch.tensor([[3, 2, 5], [3, 1, 5]]).T,
|
||||
}
|
||||
assert torch.equal(edges["A:r:B"], expected_edges["A:r:B"])
|
||||
assert torch.equal(edges["B:rr:A"], expected_edges["B:rr:A"])
|
||||
|
||||
|
||||
def test_find_reverse_edges_bi_reverse_types():
|
||||
edges = {
|
||||
"A:r:B": torch.tensor([[1, 5], [2, 5]]).T,
|
||||
"B:rr:A": torch.tensor([[3], [3]]).T,
|
||||
}
|
||||
edges = gb.add_reverse_edges(edges, {"A:r:B": "B:rr:A", "B:rr:A": "A:r:B"})
|
||||
expected_edges = {
|
||||
"A:r:B": torch.tensor([[1, 5, 3], [2, 5, 3]]).T,
|
||||
"B:rr:A": torch.tensor([[3, 2, 5], [3, 1, 5]]).T,
|
||||
}
|
||||
assert torch.equal(edges["A:r:B"], expected_edges["A:r:B"])
|
||||
assert torch.equal(edges["B:rr:A"], expected_edges["B:rr:A"])
|
||||
|
||||
|
||||
def test_find_reverse_edges_circual_reverse_types():
|
||||
edges = {
|
||||
"A:r1:B": torch.tensor([[1, 1]]),
|
||||
"B:r2:C": torch.tensor([[2, 2]]),
|
||||
"C:r3:A": torch.tensor([[3, 3]]),
|
||||
}
|
||||
edges = gb.add_reverse_edges(
|
||||
edges, {"A:r1:B": "B:r2:C", "B:r2:C": "C:r3:A", "C:r3:A": "A:r1:B"}
|
||||
)
|
||||
expected_edges = {
|
||||
"A:r1:B": torch.tensor([[1, 3], [1, 3]]).T,
|
||||
"B:r2:C": torch.tensor([[2, 1], [2, 1]]).T,
|
||||
"C:r3:A": torch.tensor([[3, 2], [3, 2]]).T,
|
||||
}
|
||||
assert torch.equal(edges["A:r1:B"], expected_edges["A:r1:B"])
|
||||
assert torch.equal(edges["B:r2:C"], expected_edges["B:r2:C"])
|
||||
assert torch.equal(edges["A:r1:B"], expected_edges["A:r1:B"])
|
||||
assert torch.equal(edges["C:r3:A"], expected_edges["C:r3:A"])
|
||||
@@ -0,0 +1,360 @@
|
||||
import dgl
|
||||
import dgl.graphbolt as gb
|
||||
import dgl.sparse as dglsp
|
||||
import torch
|
||||
|
||||
|
||||
def test_integration_link_prediction():
|
||||
torch.manual_seed(926)
|
||||
|
||||
indptr = torch.tensor([0, 0, 1, 3, 6, 8, 10])
|
||||
indices = torch.tensor([5, 3, 3, 3, 3, 4, 4, 0, 5, 4])
|
||||
|
||||
matrix_a = dglsp.from_csc(indptr, indices)
|
||||
seeds = torch.t(torch.stack(matrix_a.coo()))
|
||||
node_feature_data = torch.tensor(
|
||||
[
|
||||
[0.9634, 0.2294],
|
||||
[0.6172, 0.7865],
|
||||
[0.2109, 0.1089],
|
||||
[0.8672, 0.2276],
|
||||
[0.5503, 0.8223],
|
||||
[0.5160, 0.2486],
|
||||
]
|
||||
)
|
||||
edge_feature_data = torch.tensor(
|
||||
[
|
||||
[0.5123, 0.1709, 0.6150],
|
||||
[0.1476, 0.1902, 0.1314],
|
||||
[0.2582, 0.5203, 0.6228],
|
||||
[0.3708, 0.7631, 0.2683],
|
||||
[0.2126, 0.7878, 0.7225],
|
||||
[0.7885, 0.3414, 0.5485],
|
||||
[0.4088, 0.8200, 0.1851],
|
||||
[0.0056, 0.9469, 0.4432],
|
||||
[0.8972, 0.7511, 0.3617],
|
||||
[0.5773, 0.2199, 0.3366],
|
||||
]
|
||||
)
|
||||
|
||||
item_set = gb.ItemSet(seeds, names="seeds")
|
||||
graph = gb.fused_csc_sampling_graph(indptr, indices)
|
||||
|
||||
node_feature = gb.TorchBasedFeature(node_feature_data)
|
||||
edge_feature = gb.TorchBasedFeature(edge_feature_data)
|
||||
features = {
|
||||
("node", None, "feat"): node_feature,
|
||||
("edge", None, "feat"): edge_feature,
|
||||
}
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
datapipe = gb.ItemSampler(item_set, batch_size=4)
|
||||
datapipe = datapipe.sample_uniform_negative(graph, 2)
|
||||
fanouts = torch.LongTensor([1])
|
||||
datapipe = datapipe.sample_neighbor(graph, [fanouts, fanouts], replace=True)
|
||||
datapipe = datapipe.transform(gb.exclude_seed_edges)
|
||||
datapipe = datapipe.fetch_feature(
|
||||
feature_store, node_feature_keys=["feat"], edge_feature_keys=["feat"]
|
||||
)
|
||||
dataloader = gb.DataLoader(
|
||||
datapipe,
|
||||
)
|
||||
expected = [
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([[5, 1],
|
||||
[3, 2],
|
||||
[3, 2],
|
||||
[3, 3],
|
||||
[5, 2],
|
||||
[5, 1],
|
||||
[3, 4],
|
||||
[3, 3],
|
||||
[3, 5],
|
||||
[3, 2],
|
||||
[3, 0],
|
||||
[3, 4]]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 1, 1, 2, 2], dtype=torch.int32),
|
||||
indices=tensor([4, 5], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 1, 3, 2, 4, 0]),
|
||||
original_edge_ids=tensor([9, 7]),
|
||||
original_column_node_ids=tensor([5, 1, 3, 2, 4, 0]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 1, 1, 2, 2], dtype=torch.int32),
|
||||
indices=tensor([0, 5], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 1, 3, 2, 4, 0]),
|
||||
original_edge_ids=tensor([8, 7]),
|
||||
original_column_node_ids=tensor([5, 1, 3, 2, 4, 0]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.5160, 0.2486],
|
||||
[0.6172, 0.7865],
|
||||
[0.8672, 0.2276],
|
||||
[0.2109, 0.1089],
|
||||
[0.5503, 0.8223],
|
||||
[0.9634, 0.2294]])},
|
||||
labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
|
||||
input_nodes=tensor([5, 1, 3, 2, 4, 0]),
|
||||
indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
|
||||
edge_features=[{'feat': tensor([[0.5773, 0.2199, 0.3366],
|
||||
[0.0056, 0.9469, 0.4432]])},
|
||||
{'feat': tensor([[0.8972, 0.7511, 0.3617],
|
||||
[0.0056, 0.9469, 0.4432]])}],
|
||||
compacted_seeds=tensor([[0, 1],
|
||||
[2, 3],
|
||||
[2, 3],
|
||||
[2, 2],
|
||||
[0, 3],
|
||||
[0, 1],
|
||||
[2, 4],
|
||||
[2, 2],
|
||||
[2, 0],
|
||||
[2, 3],
|
||||
[2, 5],
|
||||
[2, 4]]),
|
||||
blocks=[Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2),
|
||||
Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2)],
|
||||
)"""
|
||||
),
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([[3, 3],
|
||||
[4, 3],
|
||||
[4, 4],
|
||||
[0, 4],
|
||||
[3, 4],
|
||||
[3, 5],
|
||||
[4, 1],
|
||||
[4, 4],
|
||||
[4, 4],
|
||||
[4, 5],
|
||||
[0, 1],
|
||||
[0, 3]]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 0, 1], dtype=torch.int32),
|
||||
indices=tensor([3], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([3, 4, 0, 5, 1]),
|
||||
original_edge_ids=tensor([0]),
|
||||
original_column_node_ids=tensor([3, 4, 0, 5, 1]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([3, 3], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([3, 4, 0, 5, 1]),
|
||||
original_edge_ids=tensor([8, 0]),
|
||||
original_column_node_ids=tensor([3, 4, 0, 5, 1]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.8672, 0.2276],
|
||||
[0.5503, 0.8223],
|
||||
[0.9634, 0.2294],
|
||||
[0.5160, 0.2486],
|
||||
[0.6172, 0.7865]])},
|
||||
labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
|
||||
input_nodes=tensor([3, 4, 0, 5, 1]),
|
||||
indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
|
||||
edge_features=[{'feat': tensor([[0.5123, 0.1709, 0.6150]])},
|
||||
{'feat': tensor([[0.8972, 0.7511, 0.3617],
|
||||
[0.5123, 0.1709, 0.6150]])}],
|
||||
compacted_seeds=tensor([[0, 0],
|
||||
[1, 0],
|
||||
[1, 1],
|
||||
[2, 1],
|
||||
[0, 1],
|
||||
[0, 3],
|
||||
[1, 4],
|
||||
[1, 1],
|
||||
[1, 1],
|
||||
[1, 3],
|
||||
[2, 4],
|
||||
[2, 0]]),
|
||||
blocks=[Block(num_src_nodes=5, num_dst_nodes=5, num_edges=1),
|
||||
Block(num_src_nodes=5, num_dst_nodes=5, num_edges=2)],
|
||||
)"""
|
||||
),
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([[5, 5],
|
||||
[4, 5],
|
||||
[5, 5],
|
||||
[5, 5],
|
||||
[4, 0],
|
||||
[4, 0]]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1], dtype=torch.int32),
|
||||
indices=tensor([1], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 4, 0]),
|
||||
original_edge_ids=tensor([6]),
|
||||
original_column_node_ids=tensor([5, 4, 0]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1], dtype=torch.int32),
|
||||
indices=tensor([2], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 4, 0]),
|
||||
original_edge_ids=tensor([7]),
|
||||
original_column_node_ids=tensor([5, 4, 0]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.5160, 0.2486],
|
||||
[0.5503, 0.8223],
|
||||
[0.9634, 0.2294]])},
|
||||
labels=tensor([1., 1., 0., 0., 0., 0.]),
|
||||
input_nodes=tensor([5, 4, 0]),
|
||||
indexes=tensor([0, 1, 0, 0, 1, 1]),
|
||||
edge_features=[{'feat': tensor([[0.4088, 0.8200, 0.1851]])},
|
||||
{'feat': tensor([[0.0056, 0.9469, 0.4432]])}],
|
||||
compacted_seeds=tensor([[0, 0],
|
||||
[1, 0],
|
||||
[0, 0],
|
||||
[0, 0],
|
||||
[1, 2],
|
||||
[1, 2]]),
|
||||
blocks=[Block(num_src_nodes=3, num_dst_nodes=3, num_edges=1),
|
||||
Block(num_src_nodes=3, num_dst_nodes=3, num_edges=1)],
|
||||
)"""
|
||||
),
|
||||
]
|
||||
for step, data in enumerate(dataloader):
|
||||
assert expected[step] == str(data), print(step, data)
|
||||
|
||||
|
||||
def test_integration_node_classification():
|
||||
torch.manual_seed(926)
|
||||
|
||||
indptr = torch.tensor([0, 0, 1, 3, 6, 8, 10])
|
||||
indices = torch.tensor([5, 3, 3, 3, 3, 4, 4, 0, 5, 4])
|
||||
|
||||
seeds = torch.tensor([5, 1, 2, 4, 3, 0])
|
||||
node_feature_data = torch.tensor(
|
||||
[
|
||||
[0.9634, 0.2294],
|
||||
[0.6172, 0.7865],
|
||||
[0.2109, 0.1089],
|
||||
[0.8672, 0.2276],
|
||||
[0.5503, 0.8223],
|
||||
[0.5160, 0.2486],
|
||||
]
|
||||
)
|
||||
edge_feature_data = torch.tensor(
|
||||
[
|
||||
[0.5123, 0.1709, 0.6150],
|
||||
[0.1476, 0.1902, 0.1314],
|
||||
[0.2582, 0.5203, 0.6228],
|
||||
[0.3708, 0.7631, 0.2683],
|
||||
[0.2126, 0.7878, 0.7225],
|
||||
[0.7885, 0.3414, 0.5485],
|
||||
[0.4088, 0.8200, 0.1851],
|
||||
[0.0056, 0.9469, 0.4432],
|
||||
[0.8972, 0.7511, 0.3617],
|
||||
[0.5773, 0.2199, 0.3366],
|
||||
]
|
||||
)
|
||||
|
||||
item_set = gb.ItemSet(seeds, names="seeds")
|
||||
graph = gb.fused_csc_sampling_graph(indptr, indices)
|
||||
|
||||
node_feature = gb.TorchBasedFeature(node_feature_data)
|
||||
edge_feature = gb.TorchBasedFeature(edge_feature_data)
|
||||
features = {
|
||||
("node", None, "feat"): node_feature,
|
||||
("edge", None, "feat"): edge_feature,
|
||||
}
|
||||
feature_store = gb.BasicFeatureStore(features)
|
||||
datapipe = gb.ItemSampler(item_set, batch_size=2)
|
||||
fanouts = torch.LongTensor([1])
|
||||
datapipe = datapipe.sample_neighbor(graph, [fanouts, fanouts], replace=True)
|
||||
datapipe = datapipe.fetch_feature(
|
||||
feature_store, node_feature_keys=["feat"], edge_feature_keys=["feat"]
|
||||
)
|
||||
dataloader = gb.DataLoader(
|
||||
datapipe,
|
||||
)
|
||||
expected = [
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([5, 1]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([0, 0], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 1]),
|
||||
original_edge_ids=tensor([8, 0]),
|
||||
original_column_node_ids=tensor([5, 1]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([0, 0], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([5, 1]),
|
||||
original_edge_ids=tensor([8, 0]),
|
||||
original_column_node_ids=tensor([5, 1]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.5160, 0.2486],
|
||||
[0.6172, 0.7865]])},
|
||||
labels=None,
|
||||
input_nodes=tensor([5, 1]),
|
||||
indexes=None,
|
||||
edge_features=[{'feat': tensor([[0.8972, 0.7511, 0.3617],
|
||||
[0.5123, 0.1709, 0.6150]])},
|
||||
{'feat': tensor([[0.8972, 0.7511, 0.3617],
|
||||
[0.5123, 0.1709, 0.6150]])}],
|
||||
compacted_seeds=None,
|
||||
blocks=[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=2),
|
||||
Block(num_src_nodes=2, num_dst_nodes=2, num_edges=2)],
|
||||
)"""
|
||||
),
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([2, 4]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2, 3], dtype=torch.int32),
|
||||
indices=tensor([2, 1, 2], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([2, 4, 3]),
|
||||
original_edge_ids=tensor([1, 6, 3]),
|
||||
original_column_node_ids=tensor([2, 4, 3]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([2, 1], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([2, 4, 3]),
|
||||
original_edge_ids=tensor([2, 6]),
|
||||
original_column_node_ids=tensor([2, 4]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.2109, 0.1089],
|
||||
[0.5503, 0.8223],
|
||||
[0.8672, 0.2276]])},
|
||||
labels=None,
|
||||
input_nodes=tensor([2, 4, 3]),
|
||||
indexes=None,
|
||||
edge_features=[{'feat': tensor([[0.1476, 0.1902, 0.1314],
|
||||
[0.4088, 0.8200, 0.1851],
|
||||
[0.3708, 0.7631, 0.2683]])},
|
||||
{'feat': tensor([[0.2582, 0.5203, 0.6228],
|
||||
[0.4088, 0.8200, 0.1851]])}],
|
||||
compacted_seeds=None,
|
||||
blocks=[Block(num_src_nodes=3, num_dst_nodes=3, num_edges=3),
|
||||
Block(num_src_nodes=3, num_dst_nodes=2, num_edges=2)],
|
||||
)"""
|
||||
),
|
||||
str(
|
||||
"""MiniBatch(seeds=tensor([3, 0]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1], dtype=torch.int32),
|
||||
indices=tensor([0], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([3, 0]),
|
||||
original_edge_ids=tensor([3]),
|
||||
original_column_node_ids=tensor([3, 0]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1], dtype=torch.int32),
|
||||
indices=tensor([0], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([3, 0]),
|
||||
original_edge_ids=tensor([3]),
|
||||
original_column_node_ids=tensor([3, 0]),
|
||||
)],
|
||||
node_features={'feat': tensor([[0.8672, 0.2276],
|
||||
[0.9634, 0.2294]])},
|
||||
labels=None,
|
||||
input_nodes=tensor([3, 0]),
|
||||
indexes=None,
|
||||
edge_features=[{'feat': tensor([[0.3708, 0.7631, 0.2683]])},
|
||||
{'feat': tensor([[0.3708, 0.7631, 0.2683]])}],
|
||||
compacted_seeds=None,
|
||||
blocks=[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1),
|
||||
Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1)],
|
||||
)"""
|
||||
),
|
||||
]
|
||||
for step, data in enumerate(dataloader):
|
||||
assert expected[step] == str(data), print(step, data)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,660 @@
|
||||
import re
|
||||
|
||||
import dgl
|
||||
import pytest
|
||||
import torch
|
||||
from dgl import graphbolt as gb
|
||||
|
||||
|
||||
def test_ItemSet_names():
|
||||
# ItemSet with single name.
|
||||
item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
|
||||
assert item_set.names == ("seeds",)
|
||||
|
||||
# ItemSet with multiple names.
|
||||
item_set = gb.ItemSet(
|
||||
(torch.arange(0, 5), torch.arange(5, 10)),
|
||||
names=("seeds", "labels"),
|
||||
)
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
|
||||
# ItemSet without name.
|
||||
item_set = gb.ItemSet(torch.arange(0, 5))
|
||||
assert item_set.names is None
|
||||
|
||||
# Integer-initiated ItemSet with excessive names.
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape("Number of items (1) and names (2) don't match."),
|
||||
):
|
||||
_ = gb.ItemSet(5, names=("seeds", "labels"))
|
||||
|
||||
# ItemSet with mismatched items and names.
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape("Number of items (1) and names (2) don't match."),
|
||||
):
|
||||
_ = gb.ItemSet(torch.arange(0, 5), names=("seeds", "labels"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
def test_ItemSet_scalar_dtype(dtype):
|
||||
item_set = gb.ItemSet(torch.tensor(5, dtype=dtype), names="seeds")
|
||||
for i, item in enumerate(item_set):
|
||||
assert i == item
|
||||
assert item.dtype == dtype
|
||||
assert item_set[2] == torch.tensor(2, dtype=dtype)
|
||||
assert torch.equal(
|
||||
item_set[slice(1, 4, 2)], torch.arange(1, 4, 2, dtype=dtype)
|
||||
)
|
||||
|
||||
|
||||
def test_ItemSet_length():
|
||||
# Integer with valid length
|
||||
num = 10
|
||||
item_set = gb.ItemSet(num)
|
||||
assert len(item_set) == 10
|
||||
# Test __iter__() method. Same as below.
|
||||
for i, item in enumerate(item_set):
|
||||
assert i == item
|
||||
|
||||
# Single iterable with valid length.
|
||||
ids = torch.arange(0, 5)
|
||||
item_set = gb.ItemSet(ids)
|
||||
assert len(item_set) == 5
|
||||
for i, item in enumerate(item_set):
|
||||
assert i == item.item()
|
||||
|
||||
# Tuple of iterables with valid length.
|
||||
item_set = gb.ItemSet((torch.arange(0, 5), torch.arange(5, 10)))
|
||||
assert len(item_set) == 5
|
||||
for i, (item1, item2) in enumerate(item_set):
|
||||
assert i == item1.item()
|
||||
assert i + 5 == item2.item()
|
||||
|
||||
class InvalidLength:
|
||||
def __iter__(self):
|
||||
return iter([0, 1, 2])
|
||||
|
||||
# Single iterable with invalid length.
|
||||
with pytest.raises(
|
||||
TypeError, match="object of type 'InvalidLength' has no len()"
|
||||
):
|
||||
item_set = gb.ItemSet(InvalidLength())
|
||||
|
||||
# Tuple of iterables with invalid length.
|
||||
with pytest.raises(
|
||||
TypeError, match="object of type 'InvalidLength' has no len()"
|
||||
):
|
||||
item_set = gb.ItemSet((InvalidLength(), InvalidLength()))
|
||||
|
||||
|
||||
def test_ItemSet_seed_nodes():
|
||||
# Node IDs with tensor.
|
||||
item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
|
||||
assert item_set.names == ("seeds",)
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert i == item.item()
|
||||
assert i == item_set[i]
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[::2], torch.tensor([0, 2, 4]))
|
||||
# Indexing with an Iterable.
|
||||
assert torch.equal(item_set[torch.arange(0, 5)], torch.arange(0, 5))
|
||||
|
||||
# Node IDs with single integer.
|
||||
item_set = gb.ItemSet(5, names="seeds")
|
||||
assert item_set.names == ("seeds",)
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert i == item.item()
|
||||
assert i == item_set[i]
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[::2], torch.tensor([0, 2, 4]))
|
||||
assert torch.equal(item_set[torch.arange(0, 5)], torch.arange(0, 5))
|
||||
# Indexing with an integer.
|
||||
assert item_set[0] == 0
|
||||
assert item_set[-1] == 4
|
||||
# Indexing that is out of range.
|
||||
with pytest.raises(IndexError, match="ItemSet index out of range."):
|
||||
_ = item_set[5]
|
||||
with pytest.raises(IndexError, match="ItemSet index out of range."):
|
||||
_ = item_set[-10]
|
||||
# Indexing with invalid input type.
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match="ItemSet indices must be int, slice, or torch.Tensor, not <class 'float'>.",
|
||||
):
|
||||
_ = item_set[1.5]
|
||||
|
||||
|
||||
def test_ItemSet_seed_nodes_labels():
|
||||
# Node IDs and labels.
|
||||
seed_nodes = torch.arange(0, 5)
|
||||
labels = torch.randint(0, 3, (5,))
|
||||
item_set = gb.ItemSet((seed_nodes, labels), names=("seeds", "labels"))
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, (seed_node, label) in enumerate(item_set):
|
||||
assert seed_node == seed_nodes[i]
|
||||
assert label == labels[i]
|
||||
assert seed_node == item_set[i][0]
|
||||
assert label == item_set[i][1]
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:][0], seed_nodes)
|
||||
assert torch.equal(item_set[:][1], labels)
|
||||
# Indexing with an Iterable.
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][0], seed_nodes)
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
|
||||
|
||||
|
||||
def test_ItemSet_node_pairs():
|
||||
# Node pairs.
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
item_set = gb.ItemSet(node_pairs, names="seeds")
|
||||
assert item_set.names == ("seeds",)
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, (src, dst) in enumerate(item_set):
|
||||
assert node_pairs[i][0] == src
|
||||
assert node_pairs[i][1] == dst
|
||||
assert node_pairs[i][0] == item_set[i][0]
|
||||
assert node_pairs[i][1] == item_set[i][1]
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:], node_pairs)
|
||||
# Indexing with an Iterable.
|
||||
assert torch.equal(item_set[torch.arange(0, 5)], node_pairs)
|
||||
|
||||
|
||||
def test_ItemSet_node_pairs_labels():
|
||||
# Node pairs and labels
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
labels = torch.randint(0, 3, (5,))
|
||||
item_set = gb.ItemSet((node_pairs, labels), names=("seeds", "labels"))
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, (node_pair, label) in enumerate(item_set):
|
||||
assert torch.equal(node_pairs[i], node_pair)
|
||||
assert labels[i] == label
|
||||
assert torch.equal(node_pairs[i], item_set[i][0])
|
||||
assert labels[i] == item_set[i][1]
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:][0], node_pairs)
|
||||
assert torch.equal(item_set[:][1], labels)
|
||||
# Indexing with an Iterable.
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][0], node_pairs)
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
|
||||
|
||||
|
||||
def test_ItemSet_node_pairs_labels_indexes():
|
||||
# Node pairs and negative destinations.
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
labels = torch.tensor([1, 1, 0, 0, 0])
|
||||
indexes = torch.tensor([0, 1, 0, 0, 1])
|
||||
item_set = gb.ItemSet(
|
||||
(node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
|
||||
)
|
||||
assert item_set.names == ("seeds", "labels", "indexes")
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, (node_pair, label, index) in enumerate(item_set):
|
||||
assert torch.equal(node_pairs[i], node_pair)
|
||||
assert torch.equal(labels[i], label)
|
||||
assert torch.equal(indexes[i], index)
|
||||
assert torch.equal(node_pairs[i], item_set[i][0])
|
||||
assert torch.equal(labels[i], item_set[i][1])
|
||||
assert torch.equal(indexes[i], item_set[i][2])
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:][0], node_pairs)
|
||||
assert torch.equal(item_set[:][1], labels)
|
||||
assert torch.equal(item_set[:][2], indexes)
|
||||
# Indexing with an Iterable.
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][0], node_pairs)
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][1], labels)
|
||||
assert torch.equal(item_set[torch.arange(0, 5)][2], indexes)
|
||||
|
||||
|
||||
def test_ItemSet_graphs():
|
||||
# Graphs.
|
||||
graphs = [dgl.rand_graph(10, 20) for _ in range(5)]
|
||||
item_set = gb.ItemSet(graphs)
|
||||
assert item_set.names is None
|
||||
# Iterating over ItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert graphs[i] == item
|
||||
assert graphs[i] == item_set[i]
|
||||
# Indexing with a slice.
|
||||
assert item_set[:] == graphs
|
||||
|
||||
|
||||
def test_HeteroItemSet_names():
|
||||
# HeteroItemSet with single name.
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(torch.arange(0, 5), names="seeds"),
|
||||
"item": gb.ItemSet(torch.arange(5, 10), names="seeds"),
|
||||
}
|
||||
)
|
||||
assert item_set.names == ("seeds",)
|
||||
|
||||
# HeteroItemSet with multiple names.
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(
|
||||
(torch.arange(0, 5), torch.arange(5, 10)),
|
||||
names=("seeds", "labels"),
|
||||
),
|
||||
"item": gb.ItemSet(
|
||||
(torch.arange(5, 10), torch.arange(10, 15)),
|
||||
names=("seeds", "labels"),
|
||||
),
|
||||
}
|
||||
)
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
|
||||
# HeteroItemSet with no name.
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(torch.arange(0, 5)),
|
||||
"item": gb.ItemSet(torch.arange(5, 10)),
|
||||
}
|
||||
)
|
||||
assert item_set.names is None
|
||||
|
||||
# HeteroItemSet with mismatched items and names.
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape("All itemsets must have the same names."),
|
||||
):
|
||||
_ = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(
|
||||
(torch.arange(0, 5), torch.arange(5, 10)),
|
||||
names=("seeds", "labels"),
|
||||
),
|
||||
"item": gb.ItemSet((torch.arange(5, 10),), names=("seeds",)),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_HeteroItemSet_length():
|
||||
# Single iterable with valid length.
|
||||
user_ids = torch.arange(0, 5)
|
||||
item_ids = torch.arange(0, 5)
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(user_ids),
|
||||
"item": gb.ItemSet(item_ids),
|
||||
}
|
||||
)
|
||||
assert len(item_set) == len(user_ids) + len(item_ids)
|
||||
|
||||
# Tuple of iterables with valid length.
|
||||
node_pairs_like = torch.arange(0, 10).reshape(-1, 2)
|
||||
neg_dsts_like = torch.arange(10, 20).reshape(-1, 2)
|
||||
node_pairs_follow = torch.arange(0, 10).reshape(-1, 2)
|
||||
neg_dsts_follow = torch.arange(10, 20).reshape(-1, 2)
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user:like:item": gb.ItemSet((node_pairs_like, neg_dsts_like)),
|
||||
"user:follow:user": gb.ItemSet(
|
||||
(node_pairs_follow, neg_dsts_follow)
|
||||
),
|
||||
}
|
||||
)
|
||||
assert len(item_set) == node_pairs_like.size(0) + node_pairs_follow.size(0)
|
||||
|
||||
class InvalidLength:
|
||||
def __iter__(self):
|
||||
return iter([0, 1, 2])
|
||||
|
||||
# Single iterable with invalid length.
|
||||
with pytest.raises(
|
||||
TypeError, match="object of type 'InvalidLength' has no len()"
|
||||
):
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(InvalidLength()),
|
||||
"item": gb.ItemSet(InvalidLength()),
|
||||
}
|
||||
)
|
||||
|
||||
# Tuple of iterables with invalid length.
|
||||
with pytest.raises(
|
||||
TypeError, match="object of type 'InvalidLength' has no len()"
|
||||
):
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user:like:item": gb.ItemSet(
|
||||
(InvalidLength(), InvalidLength())
|
||||
),
|
||||
"user:follow:user": gb.ItemSet(
|
||||
(InvalidLength(), InvalidLength())
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_HeteroItemSet_iteration_seed_nodes():
|
||||
# Node IDs.
|
||||
user_ids = torch.arange(0, 5)
|
||||
item_ids = torch.arange(5, 10)
|
||||
ids = {
|
||||
"user": gb.ItemSet(user_ids, names="seeds"),
|
||||
"item": gb.ItemSet(item_ids, names="seeds"),
|
||||
}
|
||||
chained_ids = []
|
||||
for key, value in ids.items():
|
||||
chained_ids += [(key, v) for v in value]
|
||||
item_set = gb.HeteroItemSet(ids)
|
||||
assert item_set.names == ("seeds",)
|
||||
# Iterating over HeteroItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert len(item) == 1
|
||||
assert isinstance(item, dict)
|
||||
assert chained_ids[i][0] in item
|
||||
assert item[chained_ids[i][0]] == chained_ids[i][1]
|
||||
assert item_set[i] == item
|
||||
assert item_set[i - len(item_set)] == item
|
||||
# Indexing all with a slice.
|
||||
assert torch.equal(item_set[:]["user"], user_ids)
|
||||
assert torch.equal(item_set[:]["item"], item_ids)
|
||||
# Indexing partial with a slice.
|
||||
partial_data = item_set[:3]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["user"], user_ids[:3])
|
||||
partial_data = item_set[7:]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["item"], item_ids[2:])
|
||||
partial_data = item_set[3:8:2]
|
||||
assert len(list(partial_data.keys())) == 2
|
||||
assert torch.equal(partial_data["user"], user_ids[3:-1:2])
|
||||
assert torch.equal(partial_data["item"], item_ids[0:3:2])
|
||||
# Indexing with an iterable of int.
|
||||
partial_data = item_set[torch.tensor([1, 0, 4])]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["user"], torch.tensor([1, 0, 4]))
|
||||
partial_data = item_set[torch.tensor([9, 8, 5])]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["item"], torch.tensor([9, 8, 5]))
|
||||
partial_data = item_set[torch.tensor([8, 1, 0, 9, 7, 5])]
|
||||
assert len(list(partial_data.keys())) == 2
|
||||
assert torch.equal(partial_data["user"], torch.tensor([1, 0]))
|
||||
assert torch.equal(partial_data["item"], torch.tensor([8, 9, 7, 5]))
|
||||
|
||||
# Exception cases.
|
||||
with pytest.raises(
|
||||
AssertionError, match="Start must be smaller than stop."
|
||||
):
|
||||
_ = item_set[5:3]
|
||||
with pytest.raises(
|
||||
AssertionError, match="Start must be smaller than stop."
|
||||
):
|
||||
_ = item_set[-1:3]
|
||||
with pytest.raises(IndexError, match="HeteroItemSet index out of range."):
|
||||
_ = item_set[20]
|
||||
with pytest.raises(IndexError, match="HeteroItemSet index out of range."):
|
||||
_ = item_set[-20]
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match="HeteroItemSet indices must be int, slice, or iterable of int, not <class 'float'>.",
|
||||
):
|
||||
_ = item_set[1.5]
|
||||
|
||||
|
||||
def test_HeteroItemSet_iteration_seed_nodes_labels():
|
||||
# Node IDs and labels.
|
||||
user_ids = torch.arange(0, 5)
|
||||
user_labels = torch.randint(0, 3, (5,))
|
||||
item_ids = torch.arange(5, 10)
|
||||
item_labels = torch.randint(0, 3, (5,))
|
||||
ids_labels = {
|
||||
"user": gb.ItemSet((user_ids, user_labels), names=("seeds", "labels")),
|
||||
"item": gb.ItemSet((item_ids, item_labels), names=("seeds", "labels")),
|
||||
}
|
||||
chained_ids = []
|
||||
for key, value in ids_labels.items():
|
||||
chained_ids += [(key, v) for v in value]
|
||||
item_set = gb.HeteroItemSet(ids_labels)
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
# Iterating over HeteroItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert len(item) == 1
|
||||
assert isinstance(item, dict)
|
||||
assert chained_ids[i][0] in item
|
||||
assert item[chained_ids[i][0]] == chained_ids[i][1]
|
||||
assert item_set[i] == item
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:]["user"][0], user_ids)
|
||||
assert torch.equal(item_set[:]["user"][1], user_labels)
|
||||
assert torch.equal(item_set[:]["item"][0], item_ids)
|
||||
assert torch.equal(item_set[:]["item"][1], item_labels)
|
||||
|
||||
|
||||
def test_HeteroItemSet_iteration_node_pairs():
|
||||
# Node pairs.
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
node_pairs_dict = {
|
||||
"user:like:item": gb.ItemSet(node_pairs, names="seeds"),
|
||||
"user:follow:user": gb.ItemSet(node_pairs, names="seeds"),
|
||||
}
|
||||
expected_data = []
|
||||
for key, value in node_pairs_dict.items():
|
||||
expected_data += [(key, v) for v in value]
|
||||
item_set = gb.HeteroItemSet(node_pairs_dict)
|
||||
assert item_set.names == ("seeds",)
|
||||
# Iterating over HeteroItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert len(item) == 1
|
||||
assert isinstance(item, dict)
|
||||
assert expected_data[i][0] in item
|
||||
assert torch.equal(item[expected_data[i][0]], expected_data[i][1])
|
||||
assert item_set[i].keys() == item.keys()
|
||||
key = list(item.keys())[0]
|
||||
assert torch.equal(item_set[i][key], item[key])
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:]["user:like:item"], node_pairs)
|
||||
assert torch.equal(item_set[:]["user:follow:user"], node_pairs)
|
||||
|
||||
|
||||
def test_HeteroItemSet_iteration_node_pairs_labels():
|
||||
# Node pairs and labels
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
labels = torch.randint(0, 3, (5,))
|
||||
node_pairs_labels = {
|
||||
"user:like:item": gb.ItemSet(
|
||||
(node_pairs, labels), names=("seeds", "labels")
|
||||
),
|
||||
"user:follow:user": gb.ItemSet(
|
||||
(node_pairs, labels), names=("seeds", "labels")
|
||||
),
|
||||
}
|
||||
expected_data = []
|
||||
for key, value in node_pairs_labels.items():
|
||||
expected_data += [(key, v) for v in value]
|
||||
item_set = gb.HeteroItemSet(node_pairs_labels)
|
||||
assert item_set.names == ("seeds", "labels")
|
||||
# Iterating over HeteroItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert len(item) == 1
|
||||
assert isinstance(item, dict)
|
||||
key, value = expected_data[i]
|
||||
assert key in item
|
||||
assert torch.equal(item[key][0], value[0])
|
||||
assert item[key][1] == value[1]
|
||||
assert item_set[i].keys() == item.keys()
|
||||
key = list(item.keys())[0]
|
||||
assert torch.equal(item_set[i][key][0], item[key][0])
|
||||
assert torch.equal(item_set[i][key][1], item[key][1])
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:]["user:like:item"][0], node_pairs)
|
||||
assert torch.equal(item_set[:]["user:like:item"][1], labels)
|
||||
assert torch.equal(item_set[:]["user:follow:user"][0], node_pairs)
|
||||
assert torch.equal(item_set[:]["user:follow:user"][1], labels)
|
||||
|
||||
|
||||
def test_HeteroItemSet_iteration_node_pairs_labels_indexes():
|
||||
# Node pairs and negative destinations.
|
||||
node_pairs = torch.arange(0, 10).reshape(-1, 2)
|
||||
labels = torch.tensor([1, 1, 0, 0, 0])
|
||||
indexes = torch.tensor([0, 1, 0, 0, 1])
|
||||
node_pairs_neg_dsts = {
|
||||
"user:like:item": gb.ItemSet(
|
||||
(node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
|
||||
),
|
||||
"user:follow:user": gb.ItemSet(
|
||||
(node_pairs, labels, indexes), names=("seeds", "labels", "indexes")
|
||||
),
|
||||
}
|
||||
expected_data = []
|
||||
for key, value in node_pairs_neg_dsts.items():
|
||||
expected_data += [(key, v) for v in value]
|
||||
item_set = gb.HeteroItemSet(node_pairs_neg_dsts)
|
||||
assert item_set.names == ("seeds", "labels", "indexes")
|
||||
# Iterating over HeteroItemSet and indexing one by one.
|
||||
for i, item in enumerate(item_set):
|
||||
assert len(item) == 1
|
||||
assert isinstance(item, dict)
|
||||
key, value = expected_data[i]
|
||||
assert key in item
|
||||
assert torch.equal(item[key][0], value[0])
|
||||
assert torch.equal(item[key][1], value[1])
|
||||
assert torch.equal(item[key][2], value[2])
|
||||
assert item_set[i].keys() == item.keys()
|
||||
key = list(item.keys())[0]
|
||||
assert torch.equal(item_set[i][key][0], item[key][0])
|
||||
assert torch.equal(item_set[i][key][1], item[key][1])
|
||||
assert torch.equal(item_set[i][key][2], item[key][2])
|
||||
# Indexing with a slice.
|
||||
assert torch.equal(item_set[:]["user:like:item"][0], node_pairs)
|
||||
assert torch.equal(item_set[:]["user:like:item"][1], labels)
|
||||
assert torch.equal(item_set[:]["user:like:item"][2], indexes)
|
||||
assert torch.equal(item_set[:]["user:follow:user"][0], node_pairs)
|
||||
assert torch.equal(item_set[:]["user:follow:user"][1], labels)
|
||||
assert torch.equal(item_set[:]["user:follow:user"][2], indexes)
|
||||
|
||||
|
||||
def test_ItemSet_repr():
|
||||
# ItemSet with single name.
|
||||
item_set = gb.ItemSet(torch.arange(0, 5), names="seeds")
|
||||
expected_str = (
|
||||
"ItemSet(\n"
|
||||
" items=(tensor([0, 1, 2, 3, 4]),),\n"
|
||||
" names=('seeds',),\n"
|
||||
")"
|
||||
)
|
||||
|
||||
assert str(item_set) == expected_str, item_set
|
||||
|
||||
# ItemSet with multiple names.
|
||||
item_set = gb.ItemSet(
|
||||
(torch.arange(0, 5), torch.arange(5, 10)),
|
||||
names=("seeds", "labels"),
|
||||
)
|
||||
expected_str = (
|
||||
"ItemSet(\n"
|
||||
" items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n"
|
||||
" names=('seeds', 'labels'),\n"
|
||||
")"
|
||||
)
|
||||
assert str(item_set) == expected_str, item_set
|
||||
|
||||
|
||||
def test_HeteroItemSet_repr():
|
||||
# HeteroItemSet with single name.
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(torch.arange(0, 5), names="seeds"),
|
||||
"item": gb.ItemSet(torch.arange(5, 10), names="seeds"),
|
||||
}
|
||||
)
|
||||
expected_str = (
|
||||
"HeteroItemSet(\n"
|
||||
" itemsets={'user': ItemSet(\n"
|
||||
" items=(tensor([0, 1, 2, 3, 4]),),\n"
|
||||
" names=('seeds',),\n"
|
||||
" ), 'item': ItemSet(\n"
|
||||
" items=(tensor([5, 6, 7, 8, 9]),),\n"
|
||||
" names=('seeds',),\n"
|
||||
" )},\n"
|
||||
" names=('seeds',),\n"
|
||||
")"
|
||||
)
|
||||
assert str(item_set) == expected_str, item_set
|
||||
|
||||
# HeteroItemSet with multiple names.
|
||||
item_set = gb.HeteroItemSet(
|
||||
{
|
||||
"user": gb.ItemSet(
|
||||
(torch.arange(0, 5), torch.arange(5, 10)),
|
||||
names=("seeds", "labels"),
|
||||
),
|
||||
"item": gb.ItemSet(
|
||||
(torch.arange(5, 10), torch.arange(10, 15)),
|
||||
names=("seeds", "labels"),
|
||||
),
|
||||
}
|
||||
)
|
||||
expected_str = (
|
||||
"HeteroItemSet(\n"
|
||||
" itemsets={'user': ItemSet(\n"
|
||||
" items=(tensor([0, 1, 2, 3, 4]), tensor([5, 6, 7, 8, 9])),\n"
|
||||
" names=('seeds', 'labels'),\n"
|
||||
" ), 'item': ItemSet(\n"
|
||||
" items=(tensor([5, 6, 7, 8, 9]), tensor([10, 11, 12, 13, 14])),\n"
|
||||
" names=('seeds', 'labels'),\n"
|
||||
" )},\n"
|
||||
" names=('seeds', 'labels'),\n"
|
||||
")"
|
||||
)
|
||||
assert str(item_set) == expected_str, item_set
|
||||
|
||||
|
||||
def test_deprecation_alias():
|
||||
"""Test `ItemSetDict` as the alias for `HeteroItemSet`."""
|
||||
|
||||
user_ids = torch.arange(0, 5)
|
||||
item_ids = torch.arange(5, 10)
|
||||
ids = {
|
||||
"user": gb.ItemSet(user_ids, names="seeds"),
|
||||
"item": gb.ItemSet(item_ids, names="seeds"),
|
||||
}
|
||||
with pytest.warns(
|
||||
DeprecationWarning,
|
||||
match="ItemSetDict is deprecated and will be removed in the future. Please use HeteroItemSet instead.",
|
||||
):
|
||||
item_set_dict = gb.ItemSetDict(ids)
|
||||
hetero_item_set = gb.HeteroItemSet(ids)
|
||||
assert len(item_set_dict) == len(hetero_item_set)
|
||||
assert item_set_dict.names == hetero_item_set.names
|
||||
assert item_set_dict._keys == hetero_item_set._keys
|
||||
assert torch.equal(item_set_dict._offsets, hetero_item_set._offsets)
|
||||
assert (
|
||||
repr(item_set_dict)[len("ItemSetDict") :]
|
||||
== repr(hetero_item_set)[len("HeteroItemSet") :]
|
||||
)
|
||||
# Indexing all with a slice.
|
||||
assert torch.equal(item_set_dict[:]["user"], hetero_item_set[:]["user"])
|
||||
assert torch.equal(item_set_dict[:]["item"], hetero_item_set[:]["item"])
|
||||
# Indexing partial with a slice.
|
||||
partial_data = item_set_dict[:3]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["user"], hetero_item_set[:3]["user"])
|
||||
partial_data = item_set_dict[7:]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["item"], hetero_item_set[7:]["item"])
|
||||
partial_data = item_set_dict[3:8:2]
|
||||
assert len(list(partial_data.keys())) == 2
|
||||
assert torch.equal(partial_data["user"], hetero_item_set[3:8:2]["user"])
|
||||
assert torch.equal(partial_data["item"], hetero_item_set[3:8:2]["item"])
|
||||
# Indexing with an iterable of int.
|
||||
partial_data = item_set_dict[torch.tensor([1, 0, 4])]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["user"], hetero_item_set[1, 0, 4]["user"])
|
||||
partial_data = item_set_dict[torch.tensor([9, 8, 5])]
|
||||
assert len(list(partial_data.keys())) == 1
|
||||
assert torch.equal(partial_data["item"], hetero_item_set[9, 8, 5]["item"])
|
||||
partial_data = item_set_dict[torch.tensor([8, 1, 0, 9, 7, 5])]
|
||||
assert len(list(partial_data.keys())) == 2
|
||||
assert torch.equal(partial_data["user"], hetero_item_set[1, 0]["user"])
|
||||
assert torch.equal(
|
||||
partial_data["item"], hetero_item_set[8, 9, 7, 5]["item"]
|
||||
)
|
||||
@@ -0,0 +1,755 @@
|
||||
import dgl
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
relation = "A:r:B"
|
||||
reverse_relation = "B:rr:A"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
|
||||
def test_minibatch_representation_homo(indptr_dtype, indices_dtype):
|
||||
seeds = torch.tensor([10, 11])
|
||||
csc_formats = [
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 3, 5, 6], dtype=indptr_dtype),
|
||||
indices=torch.tensor([0, 1, 2, 2, 1, 2], dtype=indices_dtype),
|
||||
),
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 3], dtype=indptr_dtype),
|
||||
indices=torch.tensor([1, 2, 0], dtype=indices_dtype),
|
||||
),
|
||||
]
|
||||
original_column_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11]),
|
||||
]
|
||||
original_row_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11, 12]),
|
||||
]
|
||||
original_edge_ids = [
|
||||
torch.tensor([19, 20, 21, 22, 25, 30]),
|
||||
torch.tensor([10, 15, 17]),
|
||||
]
|
||||
node_features = {"x": torch.tensor([5, 0, 2, 1])}
|
||||
edge_features = [
|
||||
{"x": torch.tensor([9, 0, 1, 1, 7, 4])},
|
||||
{"x": torch.tensor([0, 2, 2])},
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
input_nodes = torch.tensor([8, 1, 6, 5, 9, 0, 2, 4])
|
||||
compacted_seeds = torch.tensor([0, 1])
|
||||
labels = torch.tensor([1.0, 2.0])
|
||||
# Test minibatch without data.
|
||||
minibatch = gb.MiniBatch()
|
||||
expect_result = str(
|
||||
"""MiniBatch(seeds=None,
|
||||
sampled_subgraphs=None,
|
||||
node_features=None,
|
||||
labels=None,
|
||||
input_nodes=None,
|
||||
indexes=None,
|
||||
edge_features=None,
|
||||
compacted_seeds=None,
|
||||
blocks=None,
|
||||
)"""
|
||||
)
|
||||
result = str(minibatch)
|
||||
assert result == expect_result, print(expect_result, result)
|
||||
# Test minibatch with all attributes.
|
||||
minibatch = gb.MiniBatch(
|
||||
seeds=seeds,
|
||||
sampled_subgraphs=subgraphs,
|
||||
labels=labels,
|
||||
node_features=node_features,
|
||||
edge_features=edge_features,
|
||||
compacted_seeds=compacted_seeds,
|
||||
input_nodes=input_nodes,
|
||||
)
|
||||
expect_result = str(
|
||||
"""MiniBatch(seeds=tensor([10, 11]),
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 3, 5, 6], dtype=torch.int32),
|
||||
indices=tensor([0, 1, 2, 2, 1, 2], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([10, 11, 12, 13]),
|
||||
original_edge_ids=tensor([19, 20, 21, 22, 25, 30]),
|
||||
original_column_node_ids=tensor([10, 11, 12, 13]),
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 2, 3], dtype=torch.int32),
|
||||
indices=tensor([1, 2, 0], dtype=torch.int32),
|
||||
),
|
||||
original_row_node_ids=tensor([10, 11, 12]),
|
||||
original_edge_ids=tensor([10, 15, 17]),
|
||||
original_column_node_ids=tensor([10, 11]),
|
||||
)],
|
||||
node_features={'x': tensor([5, 0, 2, 1])},
|
||||
labels=tensor([1., 2.]),
|
||||
input_nodes=tensor([8, 1, 6, 5, 9, 0, 2, 4]),
|
||||
indexes=None,
|
||||
edge_features=[{'x': tensor([9, 0, 1, 1, 7, 4])},
|
||||
{'x': tensor([0, 2, 2])}],
|
||||
compacted_seeds=tensor([0, 1]),
|
||||
blocks=[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=6),
|
||||
Block(num_src_nodes=3, num_dst_nodes=2, num_edges=3)],
|
||||
)"""
|
||||
)
|
||||
result = str(minibatch)
|
||||
assert result == expect_result, print(expect_result, result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
|
||||
def test_minibatch_representation_hetero(indptr_dtype, indices_dtype):
|
||||
seeds = {relation: torch.tensor([10, 11])}
|
||||
csc_formats = [
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3], dtype=indptr_dtype),
|
||||
indices=torch.tensor([0, 1, 1], dtype=indices_dtype),
|
||||
),
|
||||
reverse_relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 0, 1, 2], dtype=indptr_dtype),
|
||||
indices=torch.tensor([1, 0], dtype=indices_dtype),
|
||||
),
|
||||
},
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2], dtype=indptr_dtype),
|
||||
indices=torch.tensor([1, 0], dtype=indices_dtype),
|
||||
),
|
||||
reverse_relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2], dtype=indptr_dtype),
|
||||
indices=torch.tensor([1, 0], dtype=indices_dtype),
|
||||
),
|
||||
},
|
||||
]
|
||||
original_column_node_ids = [
|
||||
{"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
|
||||
{"B": torch.tensor([10, 11]), "A": torch.tensor([5])},
|
||||
]
|
||||
original_row_node_ids = [
|
||||
{
|
||||
"A": torch.tensor([5, 7, 9, 11]),
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
},
|
||||
{
|
||||
"A": torch.tensor([5, 7]),
|
||||
"B": torch.tensor([10, 11]),
|
||||
},
|
||||
]
|
||||
original_edge_ids = [
|
||||
{
|
||||
relation: torch.tensor([19, 20, 21]),
|
||||
reverse_relation: torch.tensor([23, 26]),
|
||||
},
|
||||
{relation: torch.tensor([10, 12])},
|
||||
]
|
||||
node_features = {
|
||||
("A", "x"): torch.tensor([6, 4, 0, 1]),
|
||||
}
|
||||
edge_features = [
|
||||
{(relation, "x"): torch.tensor([4, 2, 4])},
|
||||
{(relation, "x"): torch.tensor([0, 6])},
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
compacted_seeds = {relation: torch.tensor([0, 1])}
|
||||
# Test minibatch with all attributes.
|
||||
minibatch = gb.MiniBatch(
|
||||
seeds=seeds,
|
||||
sampled_subgraphs=subgraphs,
|
||||
node_features=node_features,
|
||||
edge_features=edge_features,
|
||||
labels={"B": torch.tensor([2, 5])},
|
||||
compacted_seeds=compacted_seeds,
|
||||
input_nodes={
|
||||
"A": torch.tensor([5, 7, 9, 11]),
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
},
|
||||
)
|
||||
expect_result = str(
|
||||
"""MiniBatch(seeds={'A:r:B': tensor([10, 11])},
|
||||
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc={'A:r:B': CSCFormatBase(indptr=tensor([0, 1, 2, 3], dtype=torch.int32),
|
||||
indices=tensor([0, 1, 1], dtype=torch.int32),
|
||||
), 'B:rr:A': CSCFormatBase(indptr=tensor([0, 0, 0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([1, 0], dtype=torch.int32),
|
||||
)},
|
||||
original_row_node_ids={'A': tensor([ 5, 7, 9, 11]), 'B': tensor([10, 11, 12])},
|
||||
original_edge_ids={'A:r:B': tensor([19, 20, 21]), 'B:rr:A': tensor([23, 26])},
|
||||
original_column_node_ids={'B': tensor([10, 11, 12]), 'A': tensor([ 5, 7, 9, 11])},
|
||||
),
|
||||
SampledSubgraphImpl(sampled_csc={'A:r:B': CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
|
||||
indices=tensor([1, 0], dtype=torch.int32),
|
||||
), 'B:rr:A': CSCFormatBase(indptr=tensor([0, 2], dtype=torch.int32),
|
||||
indices=tensor([1, 0], dtype=torch.int32),
|
||||
)},
|
||||
original_row_node_ids={'A': tensor([5, 7]), 'B': tensor([10, 11])},
|
||||
original_edge_ids={'A:r:B': tensor([10, 12])},
|
||||
original_column_node_ids={'B': tensor([10, 11]), 'A': tensor([5])},
|
||||
)],
|
||||
node_features={('A', 'x'): tensor([6, 4, 0, 1])},
|
||||
labels={'B': tensor([2, 5])},
|
||||
input_nodes={'A': tensor([ 5, 7, 9, 11]), 'B': tensor([10, 11, 12])},
|
||||
indexes=None,
|
||||
edge_features=[{('A:r:B', 'x'): tensor([4, 2, 4])},
|
||||
{('A:r:B', 'x'): tensor([0, 6])}],
|
||||
compacted_seeds={'A:r:B': tensor([0, 1])},
|
||||
blocks=[Block(num_src_nodes={'A': 4, 'B': 3},
|
||||
num_dst_nodes={'A': 4, 'B': 3},
|
||||
num_edges={('A', 'r', 'B'): 3, ('B', 'rr', 'A'): 2},
|
||||
metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')]),
|
||||
Block(num_src_nodes={'A': 2, 'B': 2},
|
||||
num_dst_nodes={'A': 1, 'B': 2},
|
||||
num_edges={('A', 'r', 'B'): 2, ('B', 'rr', 'A'): 2},
|
||||
metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')])],
|
||||
)"""
|
||||
)
|
||||
result = str(minibatch)
|
||||
assert result == expect_result, print(result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("indptr_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("indices_dtype", [torch.int32, torch.int64])
|
||||
def test_get_dgl_blocks_homo(indptr_dtype, indices_dtype):
|
||||
csc_formats = [
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 3, 5, 6], dtype=indptr_dtype),
|
||||
indices=torch.tensor([0, 1, 2, 2, 1, 2], dtype=indices_dtype),
|
||||
),
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 3], dtype=indptr_dtype),
|
||||
indices=torch.tensor([0, 1, 2], dtype=indices_dtype),
|
||||
),
|
||||
]
|
||||
original_column_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11]),
|
||||
]
|
||||
original_row_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11, 12]),
|
||||
]
|
||||
original_edge_ids = [
|
||||
torch.tensor([19, 20, 21, 22, 25, 30]),
|
||||
torch.tensor([10, 15, 17]),
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
# Test minibatch with all attributes.
|
||||
minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=subgraphs,
|
||||
)
|
||||
dgl_blocks = minibatch.blocks
|
||||
expect_result = str(
|
||||
"""[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=6), Block(num_src_nodes=3, num_dst_nodes=2, num_edges=3)]"""
|
||||
)
|
||||
result = str(dgl_blocks)
|
||||
assert result == expect_result
|
||||
|
||||
|
||||
def test_get_dgl_blocks_hetero():
|
||||
csc_formats = [
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3]),
|
||||
indices=torch.tensor([0, 1, 1]),
|
||||
),
|
||||
reverse_relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 0, 1, 2]),
|
||||
indices=torch.tensor([1, 0]),
|
||||
),
|
||||
},
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2]), indices=torch.tensor([1, 0])
|
||||
),
|
||||
reverse_relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1]),
|
||||
indices=torch.tensor([1]),
|
||||
),
|
||||
},
|
||||
]
|
||||
original_column_node_ids = [
|
||||
{"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
|
||||
{"B": torch.tensor([10, 11]), "A": torch.tensor([5])},
|
||||
]
|
||||
original_row_node_ids = [
|
||||
{
|
||||
"A": torch.tensor([5, 7, 9, 11]),
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
},
|
||||
{
|
||||
"A": torch.tensor([5, 7]),
|
||||
"B": torch.tensor([10, 11]),
|
||||
},
|
||||
]
|
||||
original_edge_ids = [
|
||||
{
|
||||
relation: torch.tensor([19, 20, 21]),
|
||||
reverse_relation: torch.tensor([23, 26]),
|
||||
},
|
||||
{relation: torch.tensor([10, 12])},
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
# Test minibatch with all attributes.
|
||||
minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=subgraphs,
|
||||
)
|
||||
dgl_blocks = minibatch.blocks
|
||||
expect_result = str(
|
||||
"""[Block(num_src_nodes={'A': 4, 'B': 3},
|
||||
num_dst_nodes={'A': 4, 'B': 3},
|
||||
num_edges={('A', 'r', 'B'): 3, ('B', 'rr', 'A'): 2},
|
||||
metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')]), Block(num_src_nodes={'A': 2, 'B': 2},
|
||||
num_dst_nodes={'A': 1, 'B': 2},
|
||||
num_edges={('A', 'r', 'B'): 2, ('B', 'rr', 'A'): 1},
|
||||
metagraph=[('A', 'B', 'r'), ('B', 'A', 'rr')])]"""
|
||||
)
|
||||
result = str(dgl_blocks)
|
||||
assert result == expect_result
|
||||
|
||||
|
||||
def test_get_dgl_blocks_hetero_partial_empty_edges():
|
||||
hg = dgl.heterograph(
|
||||
{
|
||||
("n1", "e1", "n1"): ([0, 1, 1], [1, 2, 0]),
|
||||
("n1", "e2", "n2"): ([0, 1, 2], [1, 0, 2]),
|
||||
}
|
||||
)
|
||||
|
||||
gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
|
||||
|
||||
train_set = gb.HeteroItemSet(
|
||||
{"n1:e2:n2": gb.ItemSet(torch.LongTensor([[0, 1]]), names="seeds")}
|
||||
)
|
||||
datapipe = gb.ItemSampler(train_set, batch_size=1)
|
||||
datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
blocks_str = str(next(iter(dataloader)).blocks)
|
||||
expected_str = """[Block(num_src_nodes={'n1': 2, 'n2': 0},
|
||||
num_dst_nodes={'n1': 2, 'n2': 0},
|
||||
num_edges={('n1', 'e1', 'n1'): 2, ('n1', 'e2', 'n2'): 0},
|
||||
metagraph=[('n1', 'n1', 'e1'), ('n1', 'n2', 'e2')]), Block(num_src_nodes={'n1': 2, 'n2': 0},
|
||||
num_dst_nodes={'n1': 1, 'n2': 1},
|
||||
num_edges={('n1', 'e1', 'n1'): 1, ('n1', 'e2', 'n2'): 1},
|
||||
metagraph=[('n1', 'n1', 'e1'), ('n1', 'n2', 'e2')])]"""
|
||||
assert expected_str == blocks_str
|
||||
|
||||
|
||||
def test_get_dgl_blocks_hetero_empty_edges():
|
||||
hg = dgl.heterograph(
|
||||
{
|
||||
("n3", "e1", "n1"): ([0, 1, 1], [1, 2, 0]),
|
||||
("n3", "e2", "n2"): ([0, 1, 2], [1, 0, 2]),
|
||||
}
|
||||
)
|
||||
|
||||
gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
|
||||
|
||||
train_set = gb.HeteroItemSet(
|
||||
{"n3:e1:n1": gb.ItemSet(torch.LongTensor([[2, 1]]), names="seeds")}
|
||||
)
|
||||
datapipe = gb.ItemSampler(train_set, batch_size=1)
|
||||
datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
blocks_str = str(next(iter(dataloader)).blocks)
|
||||
expected_str = """[Block(num_src_nodes={'n1': 0, 'n2': 0, 'n3': 2},
|
||||
num_dst_nodes={'n1': 0, 'n2': 0, 'n3': 2},
|
||||
num_edges={('n3', 'e1', 'n1'): 0, ('n3', 'e2', 'n2'): 0},
|
||||
metagraph=[('n3', 'n1', 'e1'), ('n3', 'n2', 'e2')]), Block(num_src_nodes={'n1': 0, 'n2': 0, 'n3': 2},
|
||||
num_dst_nodes={'n1': 1, 'n2': 0, 'n3': 1},
|
||||
num_edges={('n3', 'e1', 'n1'): 1, ('n3', 'e2', 'n2'): 0},
|
||||
metagraph=[('n3', 'n1', 'e1'), ('n3', 'n2', 'e2')])]"""
|
||||
assert expected_str == blocks_str
|
||||
|
||||
|
||||
def test_get_dgl_blocks_homo_empty_edges():
|
||||
g = dgl.graph(([2, 3, 4], [3, 4, 5]))
|
||||
|
||||
gb_g = gb.from_dglgraph(g, is_homogeneous=True)
|
||||
train_set = gb.ItemSet(torch.LongTensor([[0, 1]]), names="seeds")
|
||||
datapipe = gb.ItemSampler(train_set, batch_size=1)
|
||||
datapipe = datapipe.sample_neighbor(gb_g, fanouts=[-1, -1])
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
blocks_str = str(next(iter(dataloader)).blocks)
|
||||
expected_str = "[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=0), Block(num_src_nodes=2, num_dst_nodes=2, num_edges=0)]"
|
||||
assert expected_str == blocks_str
|
||||
|
||||
|
||||
def test_seeds_ntype_being_passed():
|
||||
hg = dgl.heterograph({("n1", "e1", "n2"): ([0, 1, 2], [2, 0, 1])})
|
||||
|
||||
gb_g = gb.from_dglgraph(hg, is_homogeneous=False)
|
||||
train_set = gb.HeteroItemSet(
|
||||
{"n2": gb.ItemSet(torch.LongTensor([0, 1]), names="seeds")}
|
||||
)
|
||||
datapipe = gb.ItemSampler(train_set, batch_size=2)
|
||||
datapipe = datapipe.sample_neighbor(gb_g, [-1, -1, -1])
|
||||
dataloader = gb.DataLoader(datapipe)
|
||||
blocks = next(iter(dataloader)).blocks
|
||||
for block in blocks:
|
||||
assert "n2" in block.srctypes
|
||||
|
||||
|
||||
def create_homo_minibatch():
|
||||
csc_formats = [
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 3, 5, 6]),
|
||||
indices=torch.tensor([0, 1, 2, 2, 1, 2]),
|
||||
),
|
||||
gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 3]),
|
||||
indices=torch.tensor([1, 2, 0]),
|
||||
),
|
||||
]
|
||||
original_column_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11]),
|
||||
]
|
||||
original_row_node_ids = [
|
||||
torch.tensor([10, 11, 12, 13]),
|
||||
torch.tensor([10, 11, 12]),
|
||||
]
|
||||
original_edge_ids = [
|
||||
torch.tensor([19, 20, 21, 22, 25, 30]),
|
||||
torch.tensor([10, 15, 17]),
|
||||
]
|
||||
node_features = {"x": torch.randint(0, 10, (4,))}
|
||||
edge_features = [
|
||||
{"x": torch.randint(0, 10, (6,))},
|
||||
{"x": torch.randint(0, 10, (3,))},
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=csc_formats[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
return gb.MiniBatch(
|
||||
sampled_subgraphs=subgraphs,
|
||||
node_features=node_features,
|
||||
edge_features=edge_features,
|
||||
input_nodes=torch.tensor([10, 11, 12, 13]),
|
||||
)
|
||||
|
||||
|
||||
def create_hetero_minibatch():
|
||||
sampled_csc = [
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2, 3]),
|
||||
indices=torch.tensor([0, 1, 1]),
|
||||
),
|
||||
reverse_relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 0, 0, 1, 2]),
|
||||
indices=torch.tensor([1, 0]),
|
||||
),
|
||||
},
|
||||
{
|
||||
relation: gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 2]), indices=torch.tensor([1, 0])
|
||||
)
|
||||
},
|
||||
]
|
||||
original_column_node_ids = [
|
||||
{"B": torch.tensor([10, 11, 12]), "A": torch.tensor([5, 7, 9, 11])},
|
||||
{"B": torch.tensor([10, 11])},
|
||||
]
|
||||
original_row_node_ids = [
|
||||
{
|
||||
"A": torch.tensor([5, 7, 9, 11]),
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
},
|
||||
{
|
||||
"A": torch.tensor([5, 7]),
|
||||
"B": torch.tensor([10, 11]),
|
||||
},
|
||||
]
|
||||
original_edge_ids = [
|
||||
{
|
||||
relation: torch.tensor([19, 20, 21]),
|
||||
reverse_relation: torch.tensor([23, 26]),
|
||||
},
|
||||
{relation: torch.tensor([10, 12])},
|
||||
]
|
||||
node_features = {
|
||||
("A", "x"): torch.randint(0, 10, (4,)),
|
||||
}
|
||||
edge_features = [
|
||||
{(relation, "x"): torch.randint(0, 10, (3,))},
|
||||
{(relation, "x"): torch.randint(0, 10, (2,))},
|
||||
]
|
||||
subgraphs = []
|
||||
for i in range(2):
|
||||
subgraphs.append(
|
||||
gb.SampledSubgraphImpl(
|
||||
sampled_csc=sampled_csc[i],
|
||||
original_column_node_ids=original_column_node_ids[i],
|
||||
original_row_node_ids=original_row_node_ids[i],
|
||||
original_edge_ids=original_edge_ids[i],
|
||||
)
|
||||
)
|
||||
return gb.MiniBatch(
|
||||
sampled_subgraphs=subgraphs,
|
||||
node_features=node_features,
|
||||
edge_features=edge_features,
|
||||
input_nodes={
|
||||
"A": torch.tensor([5, 7, 9, 11]),
|
||||
"B": torch.tensor([10, 11, 12]),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def check_dgl_blocks_hetero(minibatch, blocks):
|
||||
etype = gb.etype_str_to_tuple(relation)
|
||||
sampled_csc = [
|
||||
subgraph.sampled_csc for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
original_edge_ids = [
|
||||
subgraph.original_edge_ids for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
original_row_node_ids = [
|
||||
subgraph.original_row_node_ids
|
||||
for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
|
||||
for i, block in enumerate(blocks):
|
||||
edges = block.edges(etype=etype)
|
||||
dst_ndoes = torch.arange(
|
||||
0, len(sampled_csc[i][relation].indptr) - 1
|
||||
).repeat_interleave(sampled_csc[i][relation].indptr.diff())
|
||||
assert torch.equal(edges[0], sampled_csc[i][relation].indices)
|
||||
assert torch.equal(edges[1], dst_ndoes)
|
||||
assert torch.equal(
|
||||
block.edges[etype].data[dgl.EID], original_edge_ids[i][relation]
|
||||
)
|
||||
edges = blocks[0].edges(etype=gb.etype_str_to_tuple(reverse_relation))
|
||||
dst_ndoes = torch.arange(
|
||||
0, len(sampled_csc[0][reverse_relation].indptr) - 1
|
||||
).repeat_interleave(sampled_csc[0][reverse_relation].indptr.diff())
|
||||
assert torch.equal(edges[0], sampled_csc[0][reverse_relation].indices)
|
||||
assert torch.equal(edges[1], dst_ndoes)
|
||||
assert torch.equal(
|
||||
blocks[0].srcdata[dgl.NID]["A"], original_row_node_ids[0]["A"]
|
||||
)
|
||||
assert torch.equal(
|
||||
blocks[0].srcdata[dgl.NID]["B"], original_row_node_ids[0]["B"]
|
||||
)
|
||||
|
||||
|
||||
def check_dgl_blocks_homo(minibatch, blocks):
|
||||
sampled_csc = [
|
||||
subgraph.sampled_csc for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
original_edge_ids = [
|
||||
subgraph.original_edge_ids for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
original_row_node_ids = [
|
||||
subgraph.original_row_node_ids
|
||||
for subgraph in minibatch.sampled_subgraphs
|
||||
]
|
||||
for i, block in enumerate(blocks):
|
||||
dst_ndoes = torch.arange(
|
||||
0, len(sampled_csc[i].indptr) - 1
|
||||
).repeat_interleave(sampled_csc[i].indptr.diff())
|
||||
assert torch.equal(block.edges()[0], sampled_csc[i].indices)
|
||||
assert torch.equal(block.edges()[1], dst_ndoes)
|
||||
assert torch.equal(block.edata[dgl.EID], original_edge_ids[i])
|
||||
assert torch.equal(blocks[0].srcdata[dgl.NID], original_row_node_ids[0])
|
||||
|
||||
|
||||
def test_dgl_node_classification_without_feature():
|
||||
# Arrange
|
||||
minibatch = create_homo_minibatch()
|
||||
minibatch.node_features = None
|
||||
minibatch.labels = None
|
||||
minibatch.seeds = torch.tensor([10, 15])
|
||||
# Act
|
||||
dgl_blocks = minibatch.blocks
|
||||
|
||||
# Assert
|
||||
assert len(dgl_blocks) == 2
|
||||
assert minibatch.node_features is None
|
||||
assert minibatch.labels is None
|
||||
check_dgl_blocks_homo(minibatch, dgl_blocks)
|
||||
|
||||
|
||||
def test_dgl_node_classification_homo():
|
||||
# Arrange
|
||||
minibatch = create_homo_minibatch()
|
||||
minibatch.seeds = torch.tensor([10, 15])
|
||||
minibatch.labels = torch.tensor([2, 5])
|
||||
# Act
|
||||
dgl_blocks = minibatch.blocks
|
||||
|
||||
# Assert
|
||||
assert len(dgl_blocks) == 2
|
||||
check_dgl_blocks_homo(minibatch, dgl_blocks)
|
||||
|
||||
|
||||
def test_dgl_node_classification_hetero():
|
||||
minibatch = create_hetero_minibatch()
|
||||
minibatch.labels = {"B": torch.tensor([2, 5])}
|
||||
minibatch.seeds = {"B": torch.tensor([10, 15])}
|
||||
# Act
|
||||
dgl_blocks = minibatch.blocks
|
||||
|
||||
# Assert
|
||||
assert len(dgl_blocks) == 2
|
||||
check_dgl_blocks_hetero(minibatch, dgl_blocks)
|
||||
|
||||
|
||||
def test_dgl_link_predication_homo():
|
||||
# Arrange
|
||||
minibatch = create_homo_minibatch()
|
||||
minibatch.compacted_seeds = (
|
||||
torch.tensor([[0, 1, 0, 0, 1, 1], [1, 0, 1, 1, 0, 0]]).T,
|
||||
)
|
||||
minibatch.labels = torch.tensor([1, 1, 0, 0, 0, 0])
|
||||
# Act
|
||||
dgl_blocks = minibatch.blocks
|
||||
|
||||
# Assert
|
||||
assert len(dgl_blocks) == 2
|
||||
check_dgl_blocks_homo(minibatch, dgl_blocks)
|
||||
|
||||
|
||||
def test_dgl_link_predication_hetero():
|
||||
# Arrange
|
||||
minibatch = create_hetero_minibatch()
|
||||
minibatch.compacted_seeds = {
|
||||
relation: (torch.tensor([[1, 1, 2, 0, 1, 2], [1, 0, 1, 1, 0, 0]]).T,),
|
||||
reverse_relation: (
|
||||
torch.tensor([[0, 1, 1, 2, 0, 2], [1, 0, 1, 1, 0, 0]]).T,
|
||||
),
|
||||
}
|
||||
minibatch.labels = {
|
||||
relation: (torch.tensor([1, 1, 0, 0, 0, 0]),),
|
||||
reverse_relation: (torch.tensor([1, 1, 0, 0, 0, 0]),),
|
||||
}
|
||||
# Act
|
||||
dgl_blocks = minibatch.blocks
|
||||
|
||||
# Assert
|
||||
assert len(dgl_blocks) == 2
|
||||
check_dgl_blocks_hetero(minibatch, dgl_blocks)
|
||||
|
||||
|
||||
def test_to_pyg_data():
|
||||
test_minibatch = create_homo_minibatch()
|
||||
test_minibatch.seeds = torch.tensor([0, 1])
|
||||
test_minibatch.labels = torch.tensor([7, 8])
|
||||
|
||||
expected_edge_index = torch.tensor(
|
||||
[[0, 0, 1, 1, 1, 2, 2, 2, 2], [0, 1, 0, 1, 2, 0, 1, 2, 3]]
|
||||
)
|
||||
expected_node_features = next(iter(test_minibatch.node_features.values()))
|
||||
expected_labels = torch.tensor([7, 8])
|
||||
expected_batch_size = 2
|
||||
expected_n_id = torch.tensor([10, 11, 12, 13])
|
||||
|
||||
pyg_data = test_minibatch.to_pyg_data()
|
||||
pyg_data.validate()
|
||||
assert torch.equal(pyg_data.edge_index, expected_edge_index)
|
||||
assert torch.equal(pyg_data.x, expected_node_features)
|
||||
assert torch.equal(pyg_data.y, expected_labels)
|
||||
assert pyg_data.batch_size == expected_batch_size
|
||||
assert torch.equal(pyg_data.n_id, expected_n_id)
|
||||
|
||||
test_minibatch.seeds = torch.tensor([[0, 1], [2, 3]])
|
||||
assert pyg_data.batch_size == expected_batch_size
|
||||
|
||||
test_minibatch.seeds = {"A": torch.tensor([0, 1])}
|
||||
assert pyg_data.batch_size == expected_batch_size
|
||||
|
||||
test_minibatch.seeds = {"A": torch.tensor([[0, 1], [2, 3]])}
|
||||
assert pyg_data.batch_size == expected_batch_size
|
||||
|
||||
subgraph = test_minibatch.sampled_subgraphs[0]
|
||||
# Test with sampled_csc as None.
|
||||
test_minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=None,
|
||||
node_features={"feat": expected_node_features},
|
||||
labels=expected_labels,
|
||||
)
|
||||
pyg_data = test_minibatch.to_pyg_data()
|
||||
assert pyg_data.edge_index is None, "Edge index should be none."
|
||||
|
||||
# Test with node_features as None.
|
||||
test_minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=[subgraph],
|
||||
node_features=None,
|
||||
labels=expected_labels,
|
||||
)
|
||||
pyg_data = test_minibatch.to_pyg_data()
|
||||
assert pyg_data.x is None, "Node features should be None."
|
||||
|
||||
# Test with labels as None.
|
||||
test_minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=[subgraph],
|
||||
node_features={"feat": expected_node_features},
|
||||
labels=None,
|
||||
)
|
||||
pyg_data = test_minibatch.to_pyg_data()
|
||||
assert pyg_data.y is None, "Labels should be None."
|
||||
|
||||
# Test with multiple features.
|
||||
test_minibatch = gb.MiniBatch(
|
||||
sampled_subgraphs=[subgraph],
|
||||
node_features={
|
||||
"feat": expected_node_features,
|
||||
"extra_feat": torch.tensor([[3], [4]]),
|
||||
},
|
||||
labels=expected_labels,
|
||||
)
|
||||
try:
|
||||
pyg_data = test_minibatch.to_pyg_data()
|
||||
assert (
|
||||
pyg_data.x is None
|
||||
), "Multiple features case should raise an error."
|
||||
except AssertionError as e:
|
||||
assert (
|
||||
str(e)
|
||||
== "`to_pyg_data` only supports single feature homogeneous graph."
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,308 @@
|
||||
import re
|
||||
import unittest
|
||||
|
||||
from functools import partial
|
||||
|
||||
import backend as F
|
||||
|
||||
import dgl
|
||||
import dgl.graphbolt as gb
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def test_add_reverse_edges_homo():
|
||||
edges = torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]]).T
|
||||
combined_edges = gb.add_reverse_edges(edges)
|
||||
assert torch.equal(
|
||||
combined_edges,
|
||||
torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7], [4, 5, 6, 7, 0, 1, 2, 3]]).T,
|
||||
)
|
||||
# Tensor with uncorrect dimensions.
|
||||
edges = torch.tensor([0, 1, 2, 3])
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Only tensor with shape N*2 is supported now, but got torch.Size([4])."
|
||||
),
|
||||
):
|
||||
gb.add_reverse_edges(edges)
|
||||
|
||||
|
||||
def test_add_reverse_edges_hetero():
|
||||
# reverse_etype doesn't exist in original etypes.
|
||||
edges = {"n1:e1:n2": torch.tensor([[0, 1, 2], [4, 5, 6]]).T}
|
||||
reverse_etype_mapping = {"n1:e1:n2": "n2:e2:n1"}
|
||||
combined_edges = gb.add_reverse_edges(edges, reverse_etype_mapping)
|
||||
assert torch.equal(
|
||||
combined_edges["n1:e1:n2"], torch.tensor([[0, 1, 2], [4, 5, 6]]).T
|
||||
)
|
||||
assert torch.equal(
|
||||
combined_edges["n2:e2:n1"], torch.tensor([[4, 5, 6], [0, 1, 2]]).T
|
||||
)
|
||||
# reverse_etype exists in original etypes.
|
||||
edges = {
|
||||
"n1:e1:n2": torch.tensor([[0, 1, 2], [4, 5, 6]]).T,
|
||||
"n2:e2:n1": torch.tensor([[7, 8, 9], [10, 11, 12]]).T,
|
||||
}
|
||||
reverse_etype_mapping = {"n1:e1:n2": "n2:e2:n1"}
|
||||
combined_edges = gb.add_reverse_edges(edges, reverse_etype_mapping)
|
||||
assert torch.equal(
|
||||
combined_edges["n1:e1:n2"], torch.tensor([[0, 1, 2], [4, 5, 6]]).T
|
||||
)
|
||||
assert torch.equal(
|
||||
combined_edges["n2:e2:n1"],
|
||||
torch.tensor([[7, 8, 9, 4, 5, 6], [10, 11, 12, 0, 1, 2]]).T,
|
||||
)
|
||||
# Tensor with uncorrect dimensions.
|
||||
edges = {
|
||||
"n1:e1:n2": torch.tensor([0, 1, 2]),
|
||||
"n2:e2:n1": torch.tensor([7, 8, 9]),
|
||||
}
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=re.escape(
|
||||
"Only tensor with shape N*2 is supported now, but got torch.Size([3])."
|
||||
),
|
||||
):
|
||||
gb.add_reverse_edges(edges, reverse_etype_mapping)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "gpu",
|
||||
reason="Fails due to different result on the GPU.",
|
||||
)
|
||||
@pytest.mark.parametrize("use_datapipe", [False, True])
|
||||
def test_exclude_seed_edges_homo_cpu(use_datapipe):
|
||||
graph = dgl.graph(([5, 0, 6, 7, 2, 2, 4], [0, 1, 2, 2, 3, 4, 4]))
|
||||
graph = gb.from_dglgraph(graph, True).to(F.ctx())
|
||||
items = torch.LongTensor([[0, 3], [4, 4]])
|
||||
names = "seeds"
|
||||
itemset = gb.ItemSet(items, names=names)
|
||||
datapipe = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
sampler = gb.NeighborSampler
|
||||
datapipe = sampler(datapipe, graph, fanouts)
|
||||
if use_datapipe:
|
||||
datapipe = datapipe.exclude_seed_edges()
|
||||
else:
|
||||
datapipe = datapipe.transform(partial(gb.exclude_seed_edges))
|
||||
original_row_node_ids = [
|
||||
torch.tensor([0, 3, 4, 5, 2, 6, 7]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
|
||||
]
|
||||
compacted_indices = [
|
||||
torch.tensor([3, 4, 4, 5, 6]).to(F.ctx()),
|
||||
torch.tensor([3, 4, 4]).to(F.ctx()),
|
||||
]
|
||||
indptr = [
|
||||
torch.tensor([0, 1, 2, 3, 3, 5]).to(F.ctx()),
|
||||
torch.tensor([0, 1, 2, 3]).to(F.ctx()),
|
||||
]
|
||||
seeds = [
|
||||
torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4]).to(F.ctx()),
|
||||
]
|
||||
for data in datapipe:
|
||||
for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
|
||||
assert torch.equal(
|
||||
sampled_subgraph.original_row_node_ids,
|
||||
original_row_node_ids[step],
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.sampled_csc.indices, compacted_indices[step]
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.sampled_csc.indptr, indptr[step]
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.original_column_node_ids, seeds[step]
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
F._default_context_str == "cpu",
|
||||
reason="Fails due to different result on the CPU.",
|
||||
)
|
||||
@pytest.mark.parametrize("use_datapipe", [False, True])
|
||||
@pytest.mark.parametrize("async_op", [False, True])
|
||||
def test_exclude_seed_edges_gpu(use_datapipe, async_op):
|
||||
graph = dgl.graph(([5, 0, 7, 7, 2, 4], [0, 1, 2, 2, 3, 4]))
|
||||
graph = gb.from_dglgraph(graph, is_homogeneous=True).to(F.ctx())
|
||||
items = torch.LongTensor([[0, 3], [4, 4]])
|
||||
names = "seeds"
|
||||
itemset = gb.ItemSet(items, names=names)
|
||||
datapipe = gb.ItemSampler(itemset, batch_size=4).copy_to(F.ctx())
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([-1]) for _ in range(num_layer)]
|
||||
sampler = gb.NeighborSampler
|
||||
datapipe = sampler(
|
||||
datapipe,
|
||||
graph,
|
||||
fanouts,
|
||||
deduplicate=True,
|
||||
)
|
||||
if use_datapipe:
|
||||
datapipe = datapipe.exclude_seed_edges(asynchronous=async_op)
|
||||
else:
|
||||
datapipe = datapipe.transform(
|
||||
partial(gb.exclude_seed_edges, async_op=async_op)
|
||||
)
|
||||
if torch.cuda.get_device_capability()[0] < 7:
|
||||
original_row_node_ids = [
|
||||
torch.tensor([0, 3, 4, 2, 5, 7]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4, 2, 5]).to(F.ctx()),
|
||||
]
|
||||
compacted_indices = [
|
||||
torch.tensor([4, 3, 5, 5]).to(F.ctx()),
|
||||
torch.tensor([4, 3]).to(F.ctx()),
|
||||
]
|
||||
indptr = [
|
||||
torch.tensor([0, 1, 2, 2, 5, 5]).to(F.ctx()),
|
||||
torch.tensor([0, 1, 2, 2]).to(F.ctx()),
|
||||
]
|
||||
seeds = [
|
||||
torch.tensor([0, 3, 4, 2, 5]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4]).to(F.ctx()),
|
||||
]
|
||||
else:
|
||||
original_row_node_ids = [
|
||||
torch.tensor([0, 3, 4, 5, 2, 7]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
|
||||
]
|
||||
compacted_indices = [
|
||||
torch.tensor([3, 4, 5, 5]).to(F.ctx()),
|
||||
torch.tensor([3, 4]).to(F.ctx()),
|
||||
]
|
||||
indptr = [
|
||||
torch.tensor([0, 1, 2, 2, 2, 4]).to(F.ctx()),
|
||||
torch.tensor([0, 1, 2, 2]).to(F.ctx()),
|
||||
]
|
||||
seeds = [
|
||||
torch.tensor([0, 3, 4, 5, 2]).to(F.ctx()),
|
||||
torch.tensor([0, 3, 4]).to(F.ctx()),
|
||||
]
|
||||
for data in datapipe:
|
||||
for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
|
||||
if async_op and not use_datapipe:
|
||||
sampled_subgraph = sampled_subgraph.wait()
|
||||
assert torch.equal(
|
||||
sampled_subgraph.original_row_node_ids,
|
||||
original_row_node_ids[step],
|
||||
)
|
||||
assert torch.equal(
|
||||
(sampled_subgraph.sampled_csc.indices), compacted_indices[step]
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.sampled_csc.indptr, indptr[step]
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.original_column_node_ids, seeds[step]
|
||||
)
|
||||
|
||||
|
||||
def get_hetero_graph():
|
||||
# COO graph:
|
||||
# [0, 0, 1, 1, 2, 2, 3, 3, 4, 4]
|
||||
# [2, 4, 2, 3, 0, 1, 1, 0, 0, 1]
|
||||
# [1, 1, 1, 1, 0, 0, 0, 0, 0] - > edge type.
|
||||
# num_nodes = 5, num_n1 = 2, num_n2 = 3
|
||||
ntypes = {"n1": 0, "n2": 1}
|
||||
etypes = {"n1:e1:n2": 0, "n2:e2:n1": 1}
|
||||
indptr = torch.LongTensor([0, 2, 4, 6, 8, 10])
|
||||
indices = torch.LongTensor([2, 4, 2, 3, 0, 1, 1, 0, 0, 1])
|
||||
type_per_edge = torch.LongTensor([1, 1, 1, 1, 0, 0, 0, 0, 0, 0])
|
||||
node_type_offset = torch.LongTensor([0, 2, 5])
|
||||
return gb.fused_csc_sampling_graph(
|
||||
indptr,
|
||||
indices,
|
||||
node_type_offset=node_type_offset,
|
||||
type_per_edge=type_per_edge,
|
||||
node_type_to_id=ntypes,
|
||||
edge_type_to_id=etypes,
|
||||
)
|
||||
|
||||
|
||||
def test_exclude_seed_edges_hetero():
|
||||
graph = get_hetero_graph().to(F.ctx())
|
||||
itemset = gb.HeteroItemSet(
|
||||
{"n1:e1:n2": gb.ItemSet(torch.tensor([[0, 1]]), names="seeds")}
|
||||
)
|
||||
item_sampler = gb.ItemSampler(itemset, batch_size=2).copy_to(F.ctx())
|
||||
num_layer = 2
|
||||
fanouts = [torch.LongTensor([2]) for _ in range(num_layer)]
|
||||
Sampler = gb.NeighborSampler
|
||||
datapipe = Sampler(
|
||||
item_sampler,
|
||||
graph,
|
||||
fanouts,
|
||||
deduplicate=True,
|
||||
)
|
||||
datapipe = datapipe.transform(partial(gb.exclude_seed_edges))
|
||||
csc_formats = [
|
||||
{
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1, 3, 5]),
|
||||
indices=torch.tensor([1, 0, 1, 0, 1]),
|
||||
),
|
||||
"n2:e2:n1": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2, 4]),
|
||||
indices=torch.tensor([1, 2, 1, 0]),
|
||||
),
|
||||
},
|
||||
{
|
||||
"n1:e1:n2": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 1]),
|
||||
indices=torch.tensor([1]),
|
||||
),
|
||||
"n2:e2:n1": gb.CSCFormatBase(
|
||||
indptr=torch.tensor([0, 2]),
|
||||
indices=torch.tensor([1, 2], dtype=torch.int64),
|
||||
),
|
||||
},
|
||||
]
|
||||
original_column_node_ids = [
|
||||
{
|
||||
"n1": torch.tensor([0, 1]),
|
||||
"n2": torch.tensor([0, 1, 2]),
|
||||
},
|
||||
{
|
||||
"n1": torch.tensor([0]),
|
||||
"n2": torch.tensor([1]),
|
||||
},
|
||||
]
|
||||
original_row_node_ids = [
|
||||
{
|
||||
"n1": torch.tensor([0, 1]),
|
||||
"n2": torch.tensor([0, 1, 2]),
|
||||
},
|
||||
{
|
||||
"n1": torch.tensor([0, 1]),
|
||||
"n2": torch.tensor([0, 1, 2]),
|
||||
},
|
||||
]
|
||||
for data in datapipe:
|
||||
for step, sampled_subgraph in enumerate(data.sampled_subgraphs):
|
||||
for ntype in ["n1", "n2"]:
|
||||
assert torch.equal(
|
||||
torch.sort(sampled_subgraph.original_row_node_ids[ntype])[
|
||||
0
|
||||
],
|
||||
original_row_node_ids[step][ntype].to(F.ctx()),
|
||||
)
|
||||
assert torch.equal(
|
||||
torch.sort(
|
||||
sampled_subgraph.original_column_node_ids[ntype]
|
||||
)[0],
|
||||
original_column_node_ids[step][ntype].to(F.ctx()),
|
||||
)
|
||||
for etype in ["n1:e1:n2", "n2:e2:n1"]:
|
||||
assert torch.equal(
|
||||
sampled_subgraph.sampled_csc[etype].indices,
|
||||
csc_formats[step][etype].indices.to(F.ctx()),
|
||||
)
|
||||
assert torch.equal(
|
||||
sampled_subgraph.sampled_csc[etype].indptr,
|
||||
csc_formats[step][etype].indptr.to(F.ctx()),
|
||||
)
|
||||
Reference in New Issue
Block a user