chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user