mirrors / dmlc--dgl

项目文件夹

dmlc--dgl
文件
0.6.x
T
Minjie Wang 22167f7203 [Refactor] Enable new kernel in all message passing APIs (#1953) 
* WIP: frame refactor

* new frame

* simple update_all builtin

* move all subgraph routines into the same file

* sddmm & spmm schedule; node & edge udf

* degree bucketing

* some tricky 0deg corner cases

* bug in frame append

* merge test_hetero_basics and test_basics

* some code rearange

* fix test_heterograph

* add mean spmm

* enable all builtin combinations

* pass gpu test

* pass pytorch tests

* wip

* fix some pt debugging codes

* fix bug in mxnet backward

* pass all mxnet utests

* passed tf tests

* docstring

* lint

* lint

* fix broadcasting bugs

* add warning and clamp for mean reducer

* add test for zero-degree mean

* address comments

* lint

* small fix
2020-08-07 15:40:25 +08:00

318 行
10 KiB
Python
原始文件 永久链接 Blame 文件历史

import numpy as np
import scipy.sparse as sp
import dgl
import dgl.function as fn
import backend as F
from test_utils import parametrize_dtype
D = 5
def generate_graph(idtype):
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
# create a graph where 0 is the source and 9 is the sink
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
# add a back flow from 9 to 0
g.add_edge(9, 0)
g.ndata.update({'f1' : F.randn((10,)), 'f2' : F.randn((10, D))})
weights = F.randn((17,))
g.edata.update({'e1': weights, 'e2': F.unsqueeze(weights, 1)})
return g
@parametrize_dtype
def test_v2v_update_all(idtype):
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph(idtype)
# update all
v1 = g.ndata[fld]
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata.update({fld : v1})
g.update_all(message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
# update all with edge weights
v1 = g.ndata[fld]
g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata.update({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
@parametrize_dtype
def test_v2v_snr(idtype):
u = F.tensor([0, 0, 0, 3, 4, 9], idtype)
v = F.tensor([1, 2, 3, 9, 9, 0], idtype)
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph(idtype)
# send and recv
v1 = g.ndata[fld]
g.send_and_recv((u, v), fn.copy_src(src=fld, out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata.update({fld : v1})
g.send_and_recv((u, v), message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
# send and recv with edge weights
v1 = g.ndata[fld]
g.send_and_recv((u, v), fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata.update({fld : v1})
g.send_and_recv((u, v), message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
@parametrize_dtype
def test_v2v_pull(idtype):
nodes = F.tensor([1, 2, 3, 9], idtype)
def _test(fld):
def message_func(edges):
return {'m' : edges.src[fld]}
def message_func_edge(edges):
if len(edges.src[fld].shape) == 1:
return {'m' : edges.src[fld] * edges.data['e1']}
else:
return {'m' : edges.src[fld] * edges.data['e2']}
def reduce_func(nodes):
return {fld : F.sum(nodes.mailbox['m'], 1)}
def apply_func(nodes):
return {fld : 2 * nodes.data[fld]}
g = generate_graph(idtype)
# send and recv
v1 = g.ndata[fld]
g.pull(nodes, fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata[fld] = v1
g.pull(nodes, message_func, reduce_func, apply_func)
v3 = g.ndata[fld]
assert F.allclose(v2, v3)
# send and recv with edge weights
v1 = g.ndata[fld]
g.pull(nodes, fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.ndata[fld]
g.ndata[fld] = v1
g.pull(nodes, message_func_edge, reduce_func, apply_func)
v4 = g.ndata[fld]
assert F.allclose(v2, v4)
# test 1d node features
_test('f1')
# test 2d node features
_test('f2')
@parametrize_dtype
def test_update_all_multi_fallback(idtype):
# create a graph with zero in degree nodes
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
g.ndata['h'] = F.randn((10, D))
g.edata['w1'] = F.randn((16,))
g.edata['w2'] = F.randn((16, D))
def _mfunc_hxw1(edges):
return {'m1' : edges.src['h'] * F.unsqueeze(edges.data['w1'], 1)}
def _mfunc_hxw2(edges):
return {'m2' : edges.src['h'] * edges.data['w2']}
def _rfunc_m1(nodes):
return {'o1' : F.sum(nodes.mailbox['m1'], 1)}
def _rfunc_m2(nodes):
return {'o2' : F.sum(nodes.mailbox['m2'], 1)}
def _rfunc_m1max(nodes):
return {'o3' : F.max(nodes.mailbox['m1'], 1)}
def _afunc(nodes):
ret = {}
for k, v in nodes.data.items():
if k.startswith('o'):
ret[k] = 2 * v
return ret
# compute ground truth
g.update_all(_mfunc_hxw1, _rfunc_m1, _afunc)
o1 = g.ndata.pop('o1')
g.update_all(_mfunc_hxw2, _rfunc_m2, _afunc)
o2 = g.ndata.pop('o2')
g.update_all(_mfunc_hxw1, _rfunc_m1max, _afunc)
o3 = g.ndata.pop('o3')
# v2v spmv
g.update_all(fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.sum(msg='m1', out='o1'),
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
# v2v fallback to e2v
g.update_all(fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.sum(msg='m2', out='o2'),
_afunc)
assert F.allclose(o2, g.ndata.pop('o2'))
@parametrize_dtype
def test_pull_multi_fallback(idtype):
# create a graph with zero in degree nodes
g = dgl.DGLGraph()
g = g.astype(idtype).to(F.ctx())
g.add_nodes(10)
for i in range(1, 9):
g.add_edge(0, i)
g.add_edge(i, 9)
g.ndata['h'] = F.randn((10, D))
g.edata['w1'] = F.randn((16,))
g.edata['w2'] = F.randn((16, D))
def _mfunc_hxw1(edges):
return {'m1' : edges.src['h'] * F.unsqueeze(edges.data['w1'], 1)}
def _mfunc_hxw2(edges):
return {'m2' : edges.src['h'] * edges.data['w2']}
def _rfunc_m1(nodes):
return {'o1' : F.sum(nodes.mailbox['m1'], 1)}
def _rfunc_m2(nodes):
return {'o2' : F.sum(nodes.mailbox['m2'], 1)}
def _rfunc_m1max(nodes):
return {'o3' : F.max(nodes.mailbox['m1'], 1)}
def _afunc(nodes):
ret = {}
for k, v in nodes.data.items():
if k.startswith('o'):
ret[k] = 2 * v
return ret
# nodes to pull
def _pull_nodes(nodes):
# compute ground truth
g.pull(nodes, _mfunc_hxw1, _rfunc_m1, _afunc)
o1 = g.ndata.pop('o1')
g.pull(nodes, _mfunc_hxw2, _rfunc_m2, _afunc)
o2 = g.ndata.pop('o2')
g.pull(nodes, _mfunc_hxw1, _rfunc_m1max, _afunc)
o3 = g.ndata.pop('o3')
# v2v spmv
g.pull(nodes, fn.src_mul_edge(src='h', edge='w1', out='m1'),
fn.sum(msg='m1', out='o1'),
_afunc)
assert F.allclose(o1, g.ndata.pop('o1'))
# v2v fallback to e2v
g.pull(nodes, fn.src_mul_edge(src='h', edge='w2', out='m2'),
fn.sum(msg='m2', out='o2'),
_afunc)
assert F.allclose(o2, g.ndata.pop('o2'))
# test#1: non-0deg nodes
nodes = [1, 2, 9]
_pull_nodes(nodes)
# test#2: 0deg nodes + non-0deg nodes
nodes = [0, 1, 2, 9]
_pull_nodes(nodes)
@parametrize_dtype
def test_spmv_3d_feat(idtype):
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * F.unsqueeze(F.unsqueeze(edges.data['h'], 1), 1)}
def sum_udf(nodes):
return {'h': F.sum(nodes.mailbox['sum'], 1)}
n = 100
p = 0.1
a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
g = dgl.DGLGraph(a)
g = g.astype(idtype).to(F.ctx())
m = g.number_of_edges()
# test#1: v2v with adj data
h = F.randn((n, 5, 5))
e = F.randn((m,))
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=fn.src_mul_edge('h', 'h', 'sum'), reduce_func=fn.sum('sum', 'h')) # 1
ans = g.ndata['h']
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=fn.sum('sum', 'h')) # 2
assert F.allclose(g.ndata['h'], ans)
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf) # 3
assert F.allclose(g.ndata['h'], ans)
# test#2: e2v
def src_mul_edge_udf(edges):
return {'sum': edges.src['h'] * edges.data['h']}
h = F.randn((n, 5, 5))
e = F.randn((m, 5, 5))
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=fn.src_mul_edge('h', 'h', 'sum'), reduce_func=fn.sum('sum', 'h')) # 1
ans = g.ndata['h']
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=fn.sum('sum', 'h')) # 2
assert F.allclose(g.ndata['h'], ans)
g.ndata['h'] = h
g.edata['h'] = e
g.update_all(message_func=src_mul_edge_udf, reduce_func=sum_udf) # 3
assert F.allclose(g.ndata['h'], ans)
if __name__ == '__main__':
test_v2v_update_all()
test_v2v_snr()
test_v2v_pull()
test_v2v_update_all_multi_fn()
test_v2v_snr_multi_fn()
test_e2v_update_all_multi_fn()
test_e2v_snr_multi_fn()
test_e2v_recv_multi_fn()
test_update_all_multi_fallback()
test_pull_multi_fallback()
test_spmv_3d_feat()
在新工单中引用 查看 Git Blame 复制永久链接