dmlc--dgl
Zihao Ye
83 行
2.9 KiB
Python
83 行
2.9 KiB
Python
import backend as F
|
|
import dgl
|
|
import networkx as nx
|
|
import dgl.utils as utils
|
|
from dgl import DGLGraph, ALL
|
|
from dgl.udf import NodeBatch, EdgeBatch
|
|
|
|
def test_node_batch():
|
|
g = dgl.DGLGraph(nx.path_graph(20))
|
|
feat = F.randn((g.number_of_nodes(), 10))
|
|
g.ndata['x'] = feat
|
|
|
|
# test all
|
|
v = ALL
|
|
n_repr = g.get_n_repr(v)
|
|
nbatch = NodeBatch(g, v, n_repr)
|
|
assert F.allclose(nbatch.data['x'], feat)
|
|
assert nbatch.mailbox is None
|
|
assert F.allclose(nbatch.nodes(), g.nodes())
|
|
assert nbatch.batch_size() == g.number_of_nodes()
|
|
assert len(nbatch) == g.number_of_nodes()
|
|
|
|
# test partial
|
|
v = utils.toindex(F.tensor([0, 3, 5, 7, 9]))
|
|
n_repr = g.get_n_repr(v)
|
|
nbatch = NodeBatch(g, v, n_repr)
|
|
assert F.allclose(nbatch.data['x'], F.gather_row(feat, F.tensor([0, 3, 5, 7, 9])))
|
|
assert nbatch.mailbox is None
|
|
assert F.allclose(nbatch.nodes(), F.tensor([0, 3, 5, 7, 9]))
|
|
assert nbatch.batch_size() == 5
|
|
assert len(nbatch) == 5
|
|
|
|
def test_edge_batch():
|
|
d = 10
|
|
g = dgl.DGLGraph(nx.path_graph(20))
|
|
nfeat = F.randn((g.number_of_nodes(), d))
|
|
efeat = F.randn((g.number_of_edges(), d))
|
|
g.ndata['x'] = nfeat
|
|
g.edata['x'] = efeat
|
|
|
|
# test all
|
|
eid = ALL
|
|
u, v, _ = g._graph.edges('eid')
|
|
|
|
src_data = g.get_n_repr(u)
|
|
edge_data = g.get_e_repr(eid)
|
|
dst_data = g.get_n_repr(v)
|
|
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
|
|
assert F.shape(ebatch.src['x'])[0] == g.number_of_edges() and\
|
|
F.shape(ebatch.src['x'])[1] == d
|
|
assert F.shape(ebatch.dst['x'])[0] == g.number_of_edges() and\
|
|
F.shape(ebatch.dst['x'])[1] == d
|
|
assert F.shape(ebatch.data['x'])[0] == g.number_of_edges() and\
|
|
F.shape(ebatch.data['x'])[1] == d
|
|
assert F.allclose(ebatch.edges()[0], u.tousertensor())
|
|
assert F.allclose(ebatch.edges()[1], v.tousertensor())
|
|
assert F.allclose(ebatch.edges()[2], F.arange(0, g.number_of_edges()))
|
|
assert ebatch.batch_size() == g.number_of_edges()
|
|
assert len(ebatch) == g.number_of_edges()
|
|
|
|
# test partial
|
|
eid = utils.toindex(F.tensor([0, 3, 5, 7, 11, 13, 15, 27]))
|
|
u, v, _ = g._graph.find_edges(eid)
|
|
src_data = g.get_n_repr(u)
|
|
edge_data = g.get_e_repr(eid)
|
|
dst_data = g.get_n_repr(v)
|
|
ebatch = EdgeBatch(g, (u, v, eid), src_data, edge_data, dst_data)
|
|
assert F.shape(ebatch.src['x'])[0] == 8 and\
|
|
F.shape(ebatch.src['x'])[1] == d
|
|
assert F.shape(ebatch.dst['x'])[0] == 8 and\
|
|
F.shape(ebatch.dst['x'])[1] == d
|
|
assert F.shape(ebatch.data['x'])[0] == 8 and\
|
|
F.shape(ebatch.data['x'])[1] == d
|
|
assert F.allclose(ebatch.edges()[0], u.tousertensor())
|
|
assert F.allclose(ebatch.edges()[1], v.tousertensor())
|
|
assert F.allclose(ebatch.edges()[2], eid.tousertensor())
|
|
assert ebatch.batch_size() == 8
|
|
assert len(ebatch) == 8
|
|
|
|
if __name__ == '__main__':
|
|
test_node_batch()
|
|
test_edge_batch()
|