dmlc--dgl
Minjie Wang
9b4d60799a
* WIP. remove graph arg in NodeBatch and EdgeBatch * refactor: use graph adapter for scheduler * WIP: recv * draft impl * stuck at bipartite * bipartite->unitgraph; support dsttype == srctype * pass test_query * pass test_query * pass test_view * test apply * pass udf message passing tests * pass quan's test using builtins * WIP: wildcard slicing * new construct methods * broken * good * add stack cross reducer * fix bug; fix mx * fix bug in csrmm2 when the CSR is not square * lint * removed FlattenedHeteroGraph class * WIP * prop nodes, prop edges, filter nodes/edges * add DGLGraph tests to heterograph. Fix several bugs * finish nx<->hetero graph conversion * create bipartite from nx * more spec on hetero/homo conversion * silly fixes * check node and edge types * repr * to api * adj APIs * inc * fix some lints and bugs * fix some lints * hetero/homo conversion * fix flatten test * more spec in hetero_from_homo and test * flatten using concat names * WIP: creators * rewrite hetero_from_homo in a more efficient way * remove useless variables * fix lint * subgraphs and typed subgraphs * lint & removed heterosubgraph class * lint x2 * disable heterograph mutation test * docstring update * add edge id for nx graph test * fix mx unittests * fix bug * try fix * fix unittest when cross_reducer is stack * fix ci * fix nx bipartite bug; docstring * fix scipy creation bug * lint * fix bug when converting heterograph from homograph * fix bug in hetero_from_homo about ntype order * trailing white * docstring fixes for add_foo and data views * docstring for relation slice * to_hetero and to_homo with feature support * lint * lint * DGLGraph compatibility * incidence matrix & docstring fixes * example string fixes * feature in hetero_from_relations * deduplication of edge types in to_hetero * fix lint * fix
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 = utils.toindex(slice(0, g.number_of_nodes()))
|
|
n_repr = g.get_n_repr(v)
|
|
nbatch = NodeBatch(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(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 = utils.toindex(slice(0, g.number_of_edges()))
|
|
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((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((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()
|