Skip to content

Commit

Permalink
fix the integration test.
Browse files Browse the repository at this point in the history
  • Loading branch information
@mfbalin
mfbalin committed Aug 22, 2024
1 parent 6edc601 commit cabf8b2
Showing 1 changed file with 93 additions and 91 deletions.
184 changes: 93 additions & 91 deletions tests/python/pytorch/graphbolt/test_integration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,19 +64,19 @@ def test_integration_link_prediction():
[3, 2],
[3, 2],
[3, 3],
[5, 0],
[5, 0],
[3, 3],
[5, 5],
[5, 2],
[3, 0],
[3, 4],
[3, 0],
[3, 0],
[3, 5],
[3, 3],
[3, 3],
[3, 4]]),
[3, 0]]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 2, 2, 2, 3], dtype=torch.int32),
indices=tensor([0, 5, 4], dtype=torch.int32),
indices=tensor([5, 5, 5], dtype=torch.int32),
),
original_row_node_ids=tensor([5, 1, 3, 2, 0, 4]),
original_edge_ids=tensor([8, 5, 7]),
original_edge_ids=tensor([9, 5, 6]),
original_column_node_ids=tensor([5, 1, 3, 2, 0, 4]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 1, 1, 1, 1, 2], dtype=torch.int32),
Expand All @@ -95,23 +95,23 @@ def test_integration_link_prediction():
labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
input_nodes=tensor([5, 1, 3, 2, 0, 4]),
indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
edge_features=[{'feat': tensor([[0.8972, 0.7511, 0.3617],
edge_features=[{'feat': tensor([[0.5773, 0.2199, 0.3366],
[0.7885, 0.3414, 0.5485],
[0.0056, 0.9469, 0.4432]])},
[0.4088, 0.8200, 0.1851]])},
{'feat': tensor([[0.5773, 0.2199, 0.3366],
[0.0056, 0.9469, 0.4432]])}],
compacted_seeds=tensor([[0, 1],
[2, 3],
[2, 3],
[2, 2],
[0, 4],
[0, 4],
[2, 2],
[0, 0],
[0, 3],
[2, 4],
[2, 5],
[2, 4],
[2, 4],
[2, 0],
[2, 2],
[2, 2],
[2, 5]]),
[2, 4]]),
blocks=[Block(num_src_nodes=6, num_dst_nodes=6, num_edges=3),
Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2)],
)"""
Expand All @@ -121,103 +121,104 @@ def test_integration_link_prediction():
[4, 3],
[4, 4],
[0, 4],
[3, 1],
[3, 4],
[3, 5],
[4, 2],
[4, 5],
[4, 4],
[4, 1],
[4, 3],
[0, 1],
[4, 1],
[4, 5],
[0, 2],
[0, 5]]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 1, 1, 2], dtype=torch.int32),
indices=tensor([4, 0], dtype=torch.int32),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 1, 2, 3], dtype=torch.int32),
indices=tensor([3, 3, 0], dtype=torch.int32),
),
original_row_node_ids=tensor([3, 4, 0, 1, 5, 2]),
original_edge_ids=tensor([0, 1]),
original_column_node_ids=tensor([3, 4, 0, 1, 5, 2]),
original_row_node_ids=tensor([3, 4, 0, 5, 1, 2]),
original_edge_ids=tensor([8, 0, 2]),
original_column_node_ids=tensor([3, 4, 0, 5, 1, 2]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 1, 2, 3], dtype=torch.int32),
indices=tensor([4, 4, 0], dtype=torch.int32),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 0, 0, 1, 2], dtype=torch.int32),
indices=tensor([3, 0], dtype=torch.int32),
),
original_row_node_ids=tensor([3, 4, 0, 1, 5, 2]),
original_edge_ids=tensor([0, 8, 1]),
original_column_node_ids=tensor([3, 4, 0, 1, 5, 2]),
original_row_node_ids=tensor([3, 4, 0, 5, 1, 2]),
original_edge_ids=tensor([0, 2]),
original_column_node_ids=tensor([3, 4, 0, 5, 1, 2]),
)],
node_features={'feat': tensor([[0.8672, 0.2276],
[0.5503, 0.8223],
[0.9634, 0.2294],
[0.6172, 0.7865],
[0.5160, 0.2486],
[0.6172, 0.7865],
[0.2109, 0.1089]])},
labels=tensor([1., 1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0.]),
input_nodes=tensor([3, 4, 0, 1, 5, 2]),
input_nodes=tensor([3, 4, 0, 5, 1, 2]),
indexes=tensor([0, 1, 2, 3, 0, 0, 1, 1, 2, 2, 3, 3]),
edge_features=[{'feat': tensor([[0.5123, 0.1709, 0.6150],
[0.1476, 0.1902, 0.1314]])},
edge_features=[{'feat': tensor([[0.8972, 0.7511, 0.3617],
[0.5123, 0.1709, 0.6150],
[0.2582, 0.5203, 0.6228]])},
{'feat': tensor([[0.5123, 0.1709, 0.6150],
[0.8972, 0.7511, 0.3617],
[0.1476, 0.1902, 0.1314]])}],
[0.2582, 0.5203, 0.6228]])}],
compacted_seeds=tensor([[0, 0],
[1, 0],
[1, 1],
[2, 1],
[0, 1],
[0, 3],
[0, 4],
[1, 5],
[1, 4],
[1, 1],
[1, 0],
[2, 3],
[2, 4]]),
blocks=[Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2),
Block(num_src_nodes=6, num_dst_nodes=6, num_edges=3)],
[1, 4],
[1, 3],
[2, 5],
[2, 3]]),
blocks=[Block(num_src_nodes=6, num_dst_nodes=6, num_edges=3),
Block(num_src_nodes=6, num_dst_nodes=6, num_edges=2)],
)"""
),
str(
"""MiniBatch(seeds=tensor([[5, 5],
[4, 5],
[5, 0],
[5, 4],
[4, 0],
[4, 1]]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1, 2], dtype=torch.int32),
indices=tensor([1, 0], dtype=torch.int32),
[5, 2],
[5, 2],
[4, 5],
[4, 4]]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 2, 3], dtype=torch.int32),
indices=tensor([4, 3, 1], dtype=torch.int32),
),
original_row_node_ids=tensor([5, 4, 0, 1]),
original_edge_ids=tensor([6, 0]),
original_column_node_ids=tensor([5, 4, 0, 1]),
original_row_node_ids=tensor([5, 4, 2, 3, 0]),
original_edge_ids=tensor([7, 1, 5]),
original_column_node_ids=tensor([5, 4, 2, 3]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 1, 1, 2], dtype=torch.int32),
indices=tensor([1, 0], dtype=torch.int32),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 0, 0, 1], dtype=torch.int32),
indices=tensor([3], dtype=torch.int32),
),
original_row_node_ids=tensor([5, 4, 0, 1]),
original_edge_ids=tensor([6, 0]),
original_column_node_ids=tensor([5, 4, 0, 1]),
original_row_node_ids=tensor([5, 4, 2, 3]),
original_edge_ids=tensor([1]),
original_column_node_ids=tensor([5, 4, 2]),
)],
node_features={'feat': tensor([[0.5160, 0.2486],
[0.5503, 0.8223],
[0.9634, 0.2294],
[0.6172, 0.7865]])},
[0.2109, 0.1089],
[0.8672, 0.2276],
[0.9634, 0.2294]])},
labels=tensor([1., 1., 0., 0., 0., 0.]),
input_nodes=tensor([5, 4, 0, 1]),
input_nodes=tensor([5, 4, 2, 3, 0]),
indexes=tensor([0, 1, 0, 0, 1, 1]),
edge_features=[{'feat': tensor([[0.4088, 0.8200, 0.1851],
[0.5123, 0.1709, 0.6150]])},
{'feat': tensor([[0.4088, 0.8200, 0.1851],
[0.5123, 0.1709, 0.6150]])}],
edge_features=[{'feat': tensor([[0.0056, 0.9469, 0.4432],
[0.1476, 0.1902, 0.1314],
[0.7885, 0.3414, 0.5485]])},
{'feat': tensor([[0.1476, 0.1902, 0.1314]])}],
compacted_seeds=tensor([[0, 0],
[1, 0],
[0, 2],
[0, 1],
[1, 2],
[1, 3]]),
blocks=[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=2),
Block(num_src_nodes=4, num_dst_nodes=4, num_edges=2)],
[0, 2],
[1, 0],
[1, 1]]),
blocks=[Block(num_src_nodes=5, num_dst_nodes=4, num_edges=3),
Block(num_src_nodes=4, num_dst_nodes=3, num_edges=1)],
)"""
),
]
for step, data in enumerate(dataloader):
assert expected[step] == str(data), print(data)
assert expected[step] == str(data), print(step, data)


def test_integration_node_classification():
Expand Down Expand Up @@ -274,18 +275,18 @@ def test_integration_node_classification():
expected = [
str(
"""MiniBatch(seeds=tensor([5, 1]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
indices=tensor([2, 0], dtype=torch.int32),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2, 3], dtype=torch.int32),
indices=tensor([0, 0, 2], dtype=torch.int32),
),
original_row_node_ids=tensor([5, 1, 4]),
original_edge_ids=tensor([9, 0]),
original_column_node_ids=tensor([5, 1]),
original_edge_ids=tensor([8, 0, 6]),
original_column_node_ids=tensor([5, 1, 4]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
indices=tensor([0, 0], dtype=torch.int32),
indices=tensor([2, 0], dtype=torch.int32),
),
original_row_node_ids=tensor([5, 1]),
original_edge_ids=tensor([8, 0]),
original_row_node_ids=tensor([5, 1, 4]),
original_edge_ids=tensor([9, 0]),
original_column_node_ids=tensor([5, 1]),
)],
node_features={'feat': tensor([[0.5160, 0.2486],
Expand All @@ -294,29 +295,30 @@ def test_integration_node_classification():
labels=None,
input_nodes=tensor([5, 1, 4]),
indexes=None,
edge_features=[{'feat': tensor([[0.5773, 0.2199, 0.3366],
[0.5123, 0.1709, 0.6150]])},
{'feat': tensor([[0.8972, 0.7511, 0.3617],
edge_features=[{'feat': tensor([[0.8972, 0.7511, 0.3617],
[0.5123, 0.1709, 0.6150],
[0.4088, 0.8200, 0.1851]])},
{'feat': tensor([[0.5773, 0.2199, 0.3366],
[0.5123, 0.1709, 0.6150]])}],
compacted_seeds=None,
blocks=[Block(num_src_nodes=3, num_dst_nodes=2, num_edges=2),
Block(num_src_nodes=2, num_dst_nodes=2, num_edges=2)],
blocks=[Block(num_src_nodes=3, num_dst_nodes=3, num_edges=3),
Block(num_src_nodes=3, num_dst_nodes=2, num_edges=2)],
)"""
),
str(
"""MiniBatch(seeds=tensor([2, 4]),
sampled_subgraphs=[SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2, 3, 3], dtype=torch.int32),
indices=tensor([2, 1, 2], dtype=torch.int32),
indices=tensor([2, 3, 2], dtype=torch.int32),
),
original_row_node_ids=tensor([2, 4, 3, 0]),
original_edge_ids=tensor([2, 6, 3]),
original_edge_ids=tensor([1, 7, 3]),
original_column_node_ids=tensor([2, 4, 3, 0]),
),
SampledSubgraphImpl(sampled_csc=CSCFormatBase(indptr=tensor([0, 1, 2], dtype=torch.int32),
indices=tensor([2, 3], dtype=torch.int32),
),
original_row_node_ids=tensor([2, 4, 3, 0]),
original_edge_ids=tensor([2, 7]),
original_edge_ids=tensor([1, 7]),
original_column_node_ids=tensor([2, 4]),
)],
node_features={'feat': tensor([[0.2109, 0.1089],
Expand All @@ -326,10 +328,10 @@ def test_integration_node_classification():
labels=None,
input_nodes=tensor([2, 4, 3, 0]),
indexes=None,
edge_features=[{'feat': tensor([[0.2582, 0.5203, 0.6228],
[0.4088, 0.8200, 0.1851],
edge_features=[{'feat': tensor([[0.1476, 0.1902, 0.1314],
[0.0056, 0.9469, 0.4432],
[0.3708, 0.7631, 0.2683]])},
{'feat': tensor([[0.2582, 0.5203, 0.6228],
{'feat': tensor([[0.1476, 0.1902, 0.1314],
[0.0056, 0.9469, 0.4432]])}],
compacted_seeds=None,
blocks=[Block(num_src_nodes=4, num_dst_nodes=4, num_edges=3),
Expand All @@ -349,7 +351,7 @@ def test_integration_node_classification():
indices=tensor([0], dtype=torch.int32),
),
original_row_node_ids=tensor([3, 0]),
original_edge_ids=tensor([4]),
original_edge_ids=tensor([3]),
original_column_node_ids=tensor([3, 0]),
)],
node_features={'feat': tensor([[0.8672, 0.2276],
Expand All @@ -358,7 +360,7 @@ def test_integration_node_classification():
input_nodes=tensor([3, 0]),
indexes=None,
edge_features=[{'feat': tensor([[0.2126, 0.7878, 0.7225]])},
{'feat': tensor([[0.2126, 0.7878, 0.7225]])}],
{'feat': tensor([[0.3708, 0.7631, 0.2683]])}],
compacted_seeds=None,
blocks=[Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1),
Block(num_src_nodes=2, num_dst_nodes=2, num_edges=1)],
Expand Down

0 comments on commit cabf8b2

Please sign in to comment.