test_batched_graph.py 4.36 KB
Newer Older
1
import dgl
Minjie Wang's avatar
Minjie Wang committed
2
import torch as th
3
import utils as U
4
5
6
7
8
9
10
11
12
13
14

def tree1():
    """Generate a tree
         0
        / \
       1   2
      / \
     3   4
    Edges are from leaves to root.
    """
    g = dgl.DGLGraph()
Minjie Wang's avatar
Minjie Wang committed
15
    g.add_nodes(5)
16
17
18
19
    g.add_edge(3, 1)
    g.add_edge(4, 1)
    g.add_edge(1, 0)
    g.add_edge(2, 0)
20
21
    g.ndata['h'] = th.Tensor([0, 1, 2, 3, 4])
    g.edata['h'] = th.randn(4, 10)
22
23
24
25
26
27
28
29
30
31
32
33
    return g

def tree2():
    """Generate a tree
         1
        / \
       4   3
      / \
     2   0
    Edges are from leaves to root.
    """
    g = dgl.DGLGraph()
Minjie Wang's avatar
Minjie Wang committed
34
    g.add_nodes(5)
35
36
37
38
    g.add_edge(2, 4)
    g.add_edge(0, 4)
    g.add_edge(4, 1)
    g.add_edge(3, 1)
39
40
    g.ndata['h'] = th.Tensor([0, 1, 2, 3, 4])
    g.edata['h'] = th.randn(4, 10)
41
42
43
44
45
46
47
    return g

def test_batch_unbatch():
    t1 = tree1()
    t2 = tree2()

    bg = dgl.batch([t1, t2])
Minjie Wang's avatar
Minjie Wang committed
48
49
50
51
52
53
54
    assert bg.number_of_nodes() == 10
    assert bg.number_of_edges() == 8
    assert bg.batch_size == 2
    assert bg.batch_num_nodes == [5, 5]
    assert bg.batch_num_edges == [4, 4]

    tt1, tt2 = dgl.unbatch(bg)
55
56
57
58
    assert U.allclose(t1.ndata['h'], tt1.ndata['h'])
    assert U.allclose(t1.edata['h'], tt1.edata['h'])
    assert U.allclose(t2.ndata['h'], tt2.ndata['h'])
    assert U.allclose(t2.edata['h'], tt2.edata['h'])
Minjie Wang's avatar
Minjie Wang committed
59
60
61
62
63
64
65
66
67
68
69
70
71

def test_batch_unbatch1():
    t1 = tree1()
    t2 = tree2()
    b1 = dgl.batch([t1, t2])
    b2 = dgl.batch([t2, b1])
    assert b2.number_of_nodes() == 15
    assert b2.number_of_edges() == 12
    assert b2.batch_size == 3
    assert b2.batch_num_nodes == [5, 5, 5]
    assert b2.batch_num_edges == [4, 4, 4]

    s1, s2, s3 = dgl.unbatch(b2)
72
73
74
75
76
77
    assert U.allclose(t2.ndata['h'], s1.ndata['h'])
    assert U.allclose(t2.edata['h'], s1.edata['h'])
    assert U.allclose(t1.ndata['h'], s2.ndata['h'])
    assert U.allclose(t1.edata['h'], s2.edata['h'])
    assert U.allclose(t2.ndata['h'], s3.ndata['h'])
    assert U.allclose(t2.edata['h'], s3.edata['h'])
78

79
def test_batch_send_then_recv():
80
81
82
83
    t1 = tree1()
    t2 = tree2()

    bg = dgl.batch([t1, t2])
84
85
    bg.register_message_func(lambda edges: {'m' : edges.src['h']})
    bg.register_reduce_func(lambda nodes: {'h' : th.sum(nodes.mailbox['m'], 1)})
Lingfan Yu's avatar
Lingfan Yu committed
86
87
    u = [3, 4, 2 + 5, 0 + 5]
    v = [1, 1, 4 + 5, 4 + 5]
88

89
    bg.send((u, v))
90
    bg.recv([1, 9]) # assuming recv takes in unique nodes
91

Lingfan Yu's avatar
Lingfan Yu committed
92
    t1, t2 = dgl.unbatch(bg)
93
94
    assert t1.ndata['h'][1] == 7
    assert t2.ndata['h'][4] == 2
95

96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
def test_batch_send_and_recv():
    t1 = tree1()
    t2 = tree2()

    bg = dgl.batch([t1, t2])
    bg.register_message_func(lambda edges: {'m' : edges.src['h']})
    bg.register_reduce_func(lambda nodes: {'h' : th.sum(nodes.mailbox['m'], 1)})
    u = [3, 4, 2 + 5, 0 + 5]
    v = [1, 1, 4 + 5, 4 + 5]

    bg.send_and_recv((u, v))

    t1, t2 = dgl.unbatch(bg)
    assert t1.ndata['h'][1] == 7
    assert t2.ndata['h'][4] == 2
111
112
113
114
115
116

def test_batch_propagate():
    t1 = tree1()
    t2 = tree2()

    bg = dgl.batch([t1, t2])
117
118
    bg.register_message_func(lambda edges: {'m' : edges.src['h']})
    bg.register_reduce_func(lambda nodes: {'h' : th.sum(nodes.mailbox['m'], 1)})
119
120
121
122
123
    # get leaves.

    order = []

    # step 1
Lingfan Yu's avatar
Lingfan Yu committed
124
125
    u = [3, 4, 2 + 5, 0 + 5]
    v = [1, 1, 4 + 5, 4 + 5]
126
127
128
    order.append((u, v))

    # step 2
Lingfan Yu's avatar
Lingfan Yu committed
129
130
    u = [1, 2, 4 + 5, 3 + 5]
    v = [0, 0, 1 + 5, 1 + 5]
131
132
    order.append((u, v))

GaiYu0's avatar
GaiYu0 committed
133
    bg.prop_edges(order)
Lingfan Yu's avatar
Lingfan Yu committed
134
    t1, t2 = dgl.unbatch(bg)
135

136
137
    assert t1.ndata['h'][0] == 9
    assert t2.ndata['h'][1] == 5
138

139
140
def test_batched_edge_ordering():
    g1 = dgl.DGLGraph()
Lingfan Yu's avatar
Lingfan Yu committed
141
142
    g1.add_nodes(6)
    g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
Minjie Wang's avatar
Minjie Wang committed
143
    e1 = th.randn(5, 10)
144
    g1.edata['h'] = e1
145
    g2 = dgl.DGLGraph()
Lingfan Yu's avatar
Lingfan Yu committed
146
147
    g2.add_nodes(6)
    g2.add_edges([0, 1 ,2 ,5, 4 ,5], [1, 2, 3, 4, 3, 0])
Minjie Wang's avatar
Minjie Wang committed
148
    e2 = th.randn(6, 10)
149
    g2.edata['h'] = e2
150
    g = dgl.batch([g1, g2])
151
152
    r1 = g.edata['h'][g.edge_id(4, 5)]
    r2 = g1.edata['h'][g1.edge_id(4, 5)]
Minjie Wang's avatar
Minjie Wang committed
153
    assert th.equal(r1, r2)
154

Lingfan Yu's avatar
Lingfan Yu committed
155
156
157
158
159
160
161
162
163
164
165
def test_batch_no_edge():
    g1 = dgl.DGLGraph()
    g1.add_nodes(6)
    g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
    g2 = dgl.DGLGraph()
    g2.add_nodes(6)
    g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
    g3 = dgl.DGLGraph()
    g3.add_nodes(1)  # no edges
    g = dgl.batch([g1, g3, g2]) # should not throw an error

166
167
if __name__ == '__main__':
    test_batch_unbatch()
Minjie Wang's avatar
Minjie Wang committed
168
    test_batch_unbatch1()
Lingfan Yu's avatar
Lingfan Yu committed
169
    test_batched_edge_ordering()
170
171
    test_batch_send_then_recv()
    test_batch_send_and_recv()
Lingfan Yu's avatar
Lingfan Yu committed
172
173
    test_batch_propagate()
    test_batch_no_edge()