test_kernel.py 13.1 KB
Newer Older
1
2
3
import dgl
import dgl.function as fn
import networkx as nx
4
import numpy as np
5
6
7
8
9
10
11
12
13
import backend as F
from itertools import product

def udf_copy_src(edges):
    return {'m': edges.src['u']}

def udf_copy_edge(edges):
    return {'m': edges.data['e']}

14
def udf_mean(nodes):
VoVAllen's avatar
VoVAllen committed
15
    return {'r2': F.mean(nodes.mailbox['m'], 1)}
16
17

def udf_sum(nodes):
VoVAllen's avatar
VoVAllen committed
18
    return {'r2': F.sum(nodes.mailbox['m'], 1)}
19
20
21
22
23
24
25
26

def udf_max(nodes):
    return {'r2': F.max(nodes.mailbox['m'], 1)}


D1 = 5
D2 = 3
D3 = 4
27
D4 = 10 # NOTE(xiang): used to dot feature vector
28
29
builtin = {'sum': fn.sum, 'max': fn.max, 'mean': fn.mean}
udf_reduce = {'sum': udf_sum, 'max': udf_max, 'mean': udf_mean}
30
31
32
fill_value = {'sum': 0, 'max': float("-inf")}


33
def generate_feature(g, broadcast='none', binary_op='none'):
34
35
36
    """Create graph with src, edge, dst feature. broadcast can be 'u',
    'e', 'v', 'none'
    """
37
    np.random.seed(31)
38
39
    nv = g.number_of_nodes()
    ne = g.number_of_edges()
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
    if binary_op == 'dot':
        if broadcast == 'e':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D2, 1, D4)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
        elif broadcast == 'u':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1, D4)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
        elif broadcast == 'v':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1, D4)))
        else:
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3, D4)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3, D4)))
57
    else:
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
        if broadcast == 'e':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D2, 1)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
        elif broadcast == 'u':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
        elif broadcast == 'v':
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D2, 1)))
        else:
            u = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
            e = F.tensor(np.random.uniform(-1, 1, (ne, D1, D2, D3)))
            v = F.tensor(np.random.uniform(-1, 1, (nv, D1, D2, D3)))
74
    return F.astype(u, F.float32), F.astype(v, F.float32), F.astype(e, F.float32)
75
76
77


def test_copy_src_reduce():
78
    def _test(red, partial):
79
        g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
80
81
82
        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
        # https://github.com/dmlc/dgl/issues/761
        g.add_edges(g.nodes(), g.nodes())
83
        g = g.to(F.ctx())
84
        hu, hv, he = generate_feature(g, 'none', 'none')
85
        if partial:
86
            nid = F.tensor(list(range(0, 100, 2)), g.idtype)
87
88
89
90
91
92

        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        with F.record_grad():
93
94
95
96
97
98
            if partial:
                g.pull(nid, fn.copy_src(src='u', out='m'),
                       builtin[red](msg='m', out='r1'))
            else:
                g.update_all(fn.copy_src(src='u', out='m'),
                             builtin[red](msg='m', out='r1'))
99
            r1 = g.ndata['r1']
VoVAllen's avatar
VoVAllen committed
100
            F.backward(F.reduce_sum(r1))
101
102
103
104
105
106
107
108
            n_grad1 = F.grad(g.ndata['u'])

        # reset grad
        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        with F.record_grad():
109
110
111
112
            if partial:
                g.pull(nid, udf_copy_src, udf_reduce[red])
            else:
                g.update_all(udf_copy_src, udf_reduce[red])
113
            r2 = g.ndata['r2']
VoVAllen's avatar
VoVAllen committed
114
            F.backward(F.reduce_sum(r2))
115
116
            n_grad2 = F.grad(g.ndata['u'])

117
118
119
120
121
122
123
124
125
        def _print_error(a, b):
            print("ERROR: Test copy_src_{} partial: {}".
                  format(red, partial))
            for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
                if not np.allclose(x, y):
                    print('@{} {} v.s. {}'.format(i, x, y))

        if not F.allclose(r1, r2):
            _print_error(r1, r2)
126
        assert F.allclose(r1, r2)
127
128
129
        if not F.allclose(n_grad1, n_grad2):
            print('node grad')
            _print_error(n_grad1, n_grad2)
130
131
        assert(F.allclose(n_grad1, n_grad2))

132
133
134
135
136
137
138
139
    _test('sum', False)
    _test('max', False)
    _test('mean', False)
    _test('sum', True)
    _test('max', True)
    _test('mean', True)


140
def test_copy_edge_reduce():
141
    def _test(red, partial):
142
        g = dgl.DGLGraph(nx.erdos_renyi_graph(100, 0.1))
143
144
        # NOTE(zihao): add self-loop to avoid zero-degree nodes.
        g.add_edges(g.nodes(), g.nodes())
145
        g = g.to(F.ctx())
146
        hu, hv, he = generate_feature(g, 'none', 'none')
147
        if partial:
148
            nid = F.tensor(list(range(0, 100, 2)), g.idtype)
149

150
151
152
153
154
        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        with F.record_grad():
155
156
157
158
159
160
            if partial:
                g.pull(nid, fn.copy_edge(edge='e', out='m'),
                       builtin[red](msg='m', out='r1'))
            else:
                g.update_all(fn.copy_edge(edge='e', out='m'),
                             builtin[red](msg='m', out='r1'))
161
            r1 = g.ndata['r1']
VoVAllen's avatar
VoVAllen committed
162
            F.backward(F.reduce_sum(r1))
163
164
165
166
167
168
169
170
            e_grad1 = F.grad(g.edata['e'])

        # reset grad
        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        with F.record_grad():
171
172
173
174
            if partial:
                g.pull(nid, udf_copy_edge, udf_reduce[red])
            else:
                g.update_all(udf_copy_edge, udf_reduce[red])
175
            r2 = g.ndata['r2']
VoVAllen's avatar
VoVAllen committed
176
            F.backward(F.reduce_sum(r2))
177
178
            e_grad2 = F.grad(g.edata['e'])

179
180
181
182
183
184
185
186
187
        def _print_error(a, b):
            print("ERROR: Test copy_edge_{} partial: {}".
                  format(red, partial))
            for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
                if not np.allclose(x, y):
                    print('@{} {} v.s. {}'.format(i, x, y))

        if not F.allclose(r1, r2):
            _print_error(r1, r2)
188
        assert F.allclose(r1, r2)
189
190
191
        if not F.allclose(e_grad1, e_grad2):
            print('edge gradient')
            _print_error(e_grad1, e_grad2)
192
193
        assert(F.allclose(e_grad1, e_grad2))

194
195
196
197
198
199
    _test('sum', False)
    _test('max', False)
    _test('mean', False)
    _test('sum', True)
    _test('max', True)
    _test('mean', True)
200
201
202


def test_all_binary_builtins():
203
204
    def _test(g, lhs, rhs, binary_op, reducer, partial, nid, broadcast='none'):
        # initialize node/edge features with uniform(-1, 1)
205
        hu, hv, he = generate_feature(g, broadcast, binary_op)
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        if binary_op == 'div':
            # op = div
            # lhs range: [-1, 1]
            # rhs range: [1, 2]
            # result range: [-1, 1]
            if rhs == 'u':
                hu = (hu + 3) / 2
            elif rhs == 'v':
                hv = (hv + 3) / 2
            elif rhs == 'e':
                he = (he + 3) / 2

        if binary_op == 'add' or binary_op == 'sub':
            # op = add, sub
            # lhs range: [-1/2, 1/2]
            # rhs range: [-1/2, 1/2]
            # result range: [-1, 1]
            hu = hu / 2
            hv = hv / 2
            he = he / 2

227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        builtin_msg_name = "{}_{}_{}".format(lhs, binary_op, rhs)
        builtin_msg = getattr(fn, builtin_msg_name)
        builtin_red = getattr(fn, reducer)

        def target_feature_switch(g, target):
            if target == "u":
                return g.ndata["u"]
            elif target == "v":
                return g.ndata["v"]
            else:
                return g.edata["e"]

        with F.record_grad():
244
245
246
247
248
            if partial:
                g.pull(nid, builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1'))
            else:
                g.update_all(builtin_msg(lhs, rhs, 'm'), builtin_red('m', 'r1'))
            r1 = g.ndata.pop('r1')
VoVAllen's avatar
VoVAllen committed
249
            F.backward(F.reduce_sum(r1))
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
            lhs_grad_1 = F.grad(target_feature_switch(g, lhs))
            rhs_grad_1 = F.grad(target_feature_switch(g, rhs))

        # reset grad
        g.ndata['u'] = F.attach_grad(F.clone(hu))
        g.ndata['v'] = F.attach_grad(F.clone(hv))
        g.edata['e'] = F.attach_grad(F.clone(he))

        def target_switch(edges, target):
            if target == "u":
                return edges.src
            elif target == "v":
                return edges.dst
            elif target == "e":
                return edges.data
            else:
                assert(0), "Unknown target {}".format(target)

        def mfunc(edges):
            op = getattr(F, binary_op)
270
271
272
273
274
275
276
277
278
            lhs_data = target_switch(edges, lhs)[lhs]
            rhs_data = target_switch(edges, rhs)[rhs]
            # NOTE(zihao): we need to do batched broadcast
            # e.g. (68, 3, 1) op (68, 5, 3, 4)
            while F.ndim(lhs_data) < F.ndim(rhs_data):
                lhs_data = F.unsqueeze(lhs_data, 1)
            while F.ndim(rhs_data) < F.ndim(lhs_data):
                rhs_data = F.unsqueeze(rhs_data, 1)
            return {"m": op(lhs_data, rhs_data)}
279
280
281
282
283
284

        def rfunc(nodes):
            op = getattr(F, reducer)
            return {"r2": op(nodes.mailbox['m'], 1)}

        with F.record_grad():
285
286
287
288
289
            if partial:
                g.pull(nid, mfunc, rfunc)
            else:
                g.update_all(mfunc, rfunc)
            r2 = g.ndata.pop('r2')
VoVAllen's avatar
VoVAllen committed
290
            F.backward(F.reduce_sum(r2), F.tensor([1.]))
291
292
293
            lhs_grad_2 = F.grad(target_feature_switch(g, lhs))
            rhs_grad_2 = F.grad(target_feature_switch(g, rhs))

294
        if reducer == 'prod':
295
296
297
            # increase tolerance for prod reducer
            # NOTE(zihao) as far as I know prod reducer has never
            # been used in any gnn models.
298
299
300
301
302
303
            rtol = 1e-2
            atol = 1e-2
        else:
            rtol = 1e-4
            atol = 1e-4

304
        def _print_error(a, b):
305
306
            print("ERROR: Test {}_{}_{}_{} broadcast: {} partial: {}".
                  format(lhs, binary_op, rhs, reducer, broadcast, partial))
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
            if lhs == 'u':
                lhs_data = hu
            elif lhs == 'v':
                lhs_data = hv
            elif lhs == 'e':
                lhs_data = he

            if rhs == 'u':
                rhs_data = hu
            elif rhs == 'v':
                rhs_data = hv
            elif rhs == 'e':
                rhs_data = he
            print("lhs", F.asnumpy(lhs_data).tolist())
            print("rhs", F.asnumpy(rhs_data).tolist())
322
            for i, (x, y) in enumerate(zip(F.asnumpy(a).flatten(), F.asnumpy(b).flatten())):
323
324
325
                if not np.allclose(x, y, rtol, atol):
                    print('@{} {} v.s. {}'.format(i, x, y))

326
        if not F.allclose(r1, r2, rtol, atol):
327
            _print_error(r1, r2)
328
        assert F.allclose(r1, r2, rtol, atol)
329
330

        if not F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol):
331
332
            print("left grad")
            _print_error(lhs_grad_1, lhs_grad_2)
333
        assert(F.allclose(lhs_grad_1, lhs_grad_2, rtol, atol))
334

335
        if not F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol):
336
337
            print("right grad")
            _print_error(rhs_grad_1, rhs_grad_2)
338
        assert(F.allclose(rhs_grad_1, rhs_grad_2, rtol, atol))
339
340
341

    g = dgl.DGLGraph()
    g.add_nodes(20)
342
343
    # NOTE(zihao): add self-loop to avoid zero-degree nodes.
    g.add_edges(g.nodes(), g.nodes())
344
345
346
347
348
349
350
351
352
    for i in range(2, 18):
        g.add_edge(0, i)
        g.add_edge(1, i)
        g.add_edge(i, 18)
        g.add_edge(i, 19)
    g.add_edge(18, 0)
    g.add_edge(18, 1)
    g.add_edge(19, 0)
    g.add_edge(19, 1)
353
354
    g = g.to(F.ctx())
    nid = F.tensor([0, 1, 4, 5, 7, 12, 14, 15, 18, 19], g.idtype)
355
    target = ["u", "v", "e"]
356

357
358
359
    for lhs, rhs in product(target, target):
        if lhs == rhs:
            continue
360
        for binary_op in ["add", "sub", "mul", "div", "dot"]:
361
            for reducer in ["sum", "max", "min", "prod", "mean"]:
362
                for broadcast in ["none", lhs, rhs]:
363
                    for partial in [False, True]:
364
365
                        _test(g, lhs, rhs, binary_op, reducer, partial, nid,
                              broadcast=broadcast)
366

367
if __name__ == '__main__':
VoVAllen's avatar
VoVAllen committed
368
369
    test_copy_src_reduce()
    test_copy_edge_reduce()
370
    test_all_binary_builtins()
371