Skip to content

GitLab

Unverified Commit 9c135fd5 authored by VoVAllen's avatar VoVAllen Committed by GitHub
Browse files

Merge pull request #4 from jermainewang/master 

Sync with latest commit
parents 9d3f299d 00add9f2
Show whitespace changes
Inline Side-by-side
Showing with 209 additions and 443 deletions
src/runtime/workspace_pool.h
View file @ 9c135fd5
......@@ -3,8 +3,8 @@
* \file workspace_pool.h
* \brief Workspace pool utility.
*/
#ifndef TVM_RUNTIME_WORKSPACE_POOL_H_
#define TVM_RUNTIME_WORKSPACE_POOL_H_
#ifndef DGL_RUNTIME_WORKSPACE_POOL_H_
#define DGL_RUNTIME_WORKSPACE_POOL_H_
#include <dgl/runtime/device_api.h>
#include <vector>
......@@ -58,4 +58,4 @@ class WorkspacePool {
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_WORKSPACE_POOL_H_
#endif // DGL_RUNTIME_WORKSPACE_POOL_H_
src/scheduler/scheduler.cc
View file @ 9c135fd5
// DGL Scheduler implementation
/*!
* Copyright (c) 2018 by Contributors
* \file scheduler/scheduler.cc
* \brief DGL Scheduler implementation
*/
#include <dgl/scheduler.h>
#include <unordered_map>
#include <vector>
#include <dgl/scheduler.h>
namespace dgl {
namespace sched {
......@@ -19,7 +22,7 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
// bkt: deg->dsts
std::unordered_map<int64_t, std::vector<int64_t>> bkt;
for (auto& it: in_edges) {
for (const auto& it : in_edges) {
bkt[it.second.size()].push_back(it.first);
}
......@@ -38,15 +41,15 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
int64_t* msec_ptr = static_cast<int64_t*>(mid_section->data);
// fill in bucketing ordering
for (auto& it: bkt) { // for each bucket
int64_t deg = it.first;
int64_t n_dst = it.second.size();
for (const auto& it : bkt) { // for each bucket
const int64_t deg = it.first;
const int64_t n_dst = it.second.size();
*deg_ptr++ = deg;
*nsec_ptr++ = n_dst;
*msec_ptr++ = deg * n_dst;
for (auto dst: it.second) { // for each dst in this bucket
for (const auto dst : it.second) { // for each dst in this bucket
*nid_ptr++ = dst;
for (auto mid: in_edges[dst]) { // for each in edge of dst
for (const auto mid : in_edges[dst]) { // for each in edge of dst
*mid_ptr++ = mid;
}
}
......
src/scheduler/scheduler_apis.cc
View file @ 9c135fd5
#include "../c_api_common.h"
/*!
* Copyright (c) 2018 by Contributors
* \file scheduler/scheduler_apis.cc
* \brief DGL scheduler APIs
*/
#include <dgl/graph.h>
#include <dgl/scheduler.h>
#include "../c_api_common.h"
using tvm::runtime::TVMArgs;
using tvm::runtime::TVMRetValue;
......@@ -18,7 +23,7 @@ TVM_REGISTER_GLOBAL("scheduler._CAPI_DGLDegreeBucketingFromGraph")
.set_body([] (TVMArgs args, TVMRetValue* rv) {
GraphHandle ghandle = args[0];
const Graph* gptr = static_cast<Graph*>(ghandle);
auto edges = gptr->Edges(false);
const auto& edges = gptr->Edges(false);
*rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(edges.dst));
});
......
tests/mxnet/test_basics.py
View file @ 9c135fd5
......@@ -9,7 +9,7 @@ reduce_msg_shapes = set()
def check_eq(a, b):
assert a.shape == b.shape
assert mx.sum(a == b) == int(np.prod(list(a.shape)))
assert mx.nd.sum(a == b).asnumpy() == int(np.prod(list(a.shape)))
def message_func(src, edge):
assert len(src['h'].shape) == 2
......@@ -53,16 +53,12 @@ def test_batch_setter_getter():
assert len(g.get_n_repr()) == 0
g.set_n_repr({'h' : mx.nd.zeros((10, D))})
# set partial nodes
# TODO we need to enable the test later.
'''
u = mx.nd.array([1, 3, 5], dtype='int64')
g.set_n_repr({'h' : mx.nd.ones((3, D))}, u)
assert _pfc(g.get_n_repr()['h']) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
# get partial nodes
u = mx.nd.array([1, 2, 3], dtype='int64')
print(g.get_n_repr(u)['h'])
assert _pfc(g.get_n_repr(u)['h']) == [1., 0., 1.]
'''
'''
s, d, eid
......@@ -127,9 +123,11 @@ def test_batch_setter_autograd():
with mx.autograd.record():
g = generate_graph(grad=True)
h1 = g.get_n_repr()['h']
h1.attach_grad()
# partial set
v = mx.nd.array([1, 2, 8], dtype='int64')
hh = mx.nd.zeros((len(v), D))
hh.attach_grad()
g.set_n_repr({'h' : hh}, v)
h2 = g.get_n_repr()['h']
h2.backward(mx.nd.ones((10, D)) * 2)
......@@ -252,8 +250,7 @@ def test_pull_0deg():
if __name__ == '__main__':
test_batch_setter_getter()
# TODO we need to enable it after index_copy is implemented.
#test_batch_setter_autograd()
test_batch_setter_autograd()
test_batch_send()
test_batch_recv()
test_update_routines()
......
tests/pytorch/test_basics.py
View file @ 9c135fd5
......@@ -20,22 +20,26 @@ def reduce_func(node, msgs):
reduce_msg_shapes.add(tuple(msgs.shape))
assert len(msgs.shape) == 3
assert msgs.shape[2] == D
return {'m' : th.sum(msgs, 1)}
return {'accum' : th.sum(msgs, 1)}
def apply_node_func(node):
return {'h' : node['h'] + node['m']}
return {'h' : node['h'] + node['accum']}
def generate_graph(grad=False):
g = DGLGraph()
g.add_nodes(10) # 10 nodes.
# create a graph where 0 is the source and 9 is the sink
# 17 edges
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)
ncol = Variable(th.randn(10, D), requires_grad=grad)
accumcol = Variable(th.randn(10, D), requires_grad=grad)
ecol = Variable(th.randn(17, D), requires_grad=grad)
g.set_n_repr({'h' : ncol})
g.set_n_initializer(lambda shape, dtype : th.zeros(shape))
return g
def test_batch_setter_getter():
......@@ -46,8 +50,9 @@ def test_batch_setter_getter():
g.set_n_repr({'h' : th.zeros((10, D))})
assert _pfc(g.get_n_repr()['h']) == [0.] * 10
# pop nodes
old_len = len(g.get_n_repr())
assert _pfc(g.pop_n_repr('h')) == [0.] * 10
assert len(g.get_n_repr()) == 0
assert len(g.get_n_repr()) == old_len - 1
g.set_n_repr({'h' : th.zeros((10, D))})
# set partial nodes
u = th.tensor([1, 3, 5])
......@@ -81,8 +86,9 @@ def test_batch_setter_getter():
g.set_e_repr({'l' : th.zeros((17, D))})
assert _pfc(g.get_e_repr()['l']) == [0.] * 17
# pop edges
old_len = len(g.get_e_repr())
assert _pfc(g.pop_e_repr('l')) == [0.] * 17
assert len(g.get_e_repr()) == 0
assert len(g.get_e_repr()) == old_len - 1
g.set_e_repr({'l' : th.zeros((17, D))})
# set partial edges (many-many)
u = th.tensor([0, 0, 2, 5, 9])
......@@ -203,14 +209,13 @@ def test_reduce_0deg():
g.add_edge(3, 0)
g.add_edge(4, 0)
def _message(src, edge):
return src
return {'m' : src['h']}
def _reduce(node, msgs):
assert msgs is not None
return node + msgs.sum(1)
return {'h' : node['h'] + msgs['m'].sum(1)}
old_repr = th.randn(5, 5)
g.set_n_repr(old_repr)
g.set_n_repr({'h' : old_repr})
g.update_all(_message, _reduce)
new_repr = g.get_n_repr()
new_repr = g.get_n_repr()['h']
assert th.allclose(new_repr[1:], old_repr[1:])
assert th.allclose(new_repr[0], old_repr.sum(0))
......@@ -220,29 +225,30 @@ def test_pull_0deg():
g.add_nodes(2)
g.add_edge(0, 1)
def _message(src, edge):
return src
return {'m' : src['h']}
def _reduce(node, msgs):
assert msgs is not None
return msgs.sum(1)
return {'h' : msgs['m'].sum(1)}
old_repr = th.randn(2, 5)
g.set_n_repr(old_repr)
g.set_n_repr({'h' : old_repr})
g.pull(0, _message, _reduce)
new_repr = g.get_n_repr()
new_repr = g.get_n_repr()['h']
assert th.allclose(new_repr[0], old_repr[0])
assert th.allclose(new_repr[1], old_repr[1])
g.pull(1, _message, _reduce)
new_repr = g.get_n_repr()
new_repr = g.get_n_repr()['h']
assert th.allclose(new_repr[1], old_repr[0])
old_repr = th.randn(2, 5)
g.set_n_repr(old_repr)
g.set_n_repr({'h' : old_repr})
g.pull([0, 1], _message, _reduce)
new_repr = g.get_n_repr()
new_repr = g.get_n_repr()['h']
assert th.allclose(new_repr[0], old_repr[0])
assert th.allclose(new_repr[1], old_repr[0])
def test_send_twice():
def _disabled_test_send_twice():
# TODO(minjie): please re-enable this unittest after the send code problem is fixed.
g = DGLGraph()
g.add_nodes(3)
g.add_edge(0, 1)
......@@ -348,5 +354,4 @@ if __name__ == '__main__':
test_update_routines()
test_reduce_0deg()
test_pull_0deg()
test_send_twice()
test_send_multigraph()
tests/pytorch/test_basics_anonymous.py deleted 100644 → 0
View file @ 9d3f299d
import torch as th
from torch.autograd import Variable
import numpy as np
from dgl.graph import DGLGraph, __REPR__
D = 32
reduce_msg_shapes = set()
def check_eq(a, b):
assert a.shape == b.shape
assert th.sum(a == b) == int(np.prod(list(a.shape)))
def message_func(hu, e_uv):
assert len(hu.shape) == 2
assert hu.shape[1] == D
return hu
def reduce_func(hv, msgs):
reduce_msg_shapes.add(tuple(msgs.shape))
assert len(msgs.shape) == 3
assert msgs.shape[2] == D
return hv + th.sum(msgs, 1)
def generate_graph(grad=False):
g = DGLGraph()
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)
col = Variable(th.randn(10, D), requires_grad=grad)
g.set_n_repr(col)
return g
def test_batch_setter_getter():
def _pfc(x):
return list(x.numpy()[:,0])
g = generate_graph()
# set all nodes
g.set_n_repr(th.zeros((10, D)))
assert _pfc(g.get_n_repr()) == [0.] * 10
# pop nodes
assert _pfc(g.pop_n_repr()) == [0.] * 10
assert len(g.get_n_repr()) == 0
g.set_n_repr(th.zeros((10, D)))
# set partial nodes
u = th.tensor([1, 3, 5])
g.set_n_repr(th.ones((3, D)), u)
assert _pfc(g.get_n_repr()) == [0., 1., 0., 1., 0., 1., 0., 0., 0., 0.]
# get partial nodes
u = th.tensor([1, 2, 3])
assert _pfc(g.get_n_repr(u)) == [1., 0., 1.]
'''
s, d, eid
0, 1, 0
1, 9, 1
0, 2, 2
2, 9, 3
0, 3, 4
3, 9, 5
0, 4, 6
4, 9, 7
0, 5, 8
5, 9, 9
0, 6, 10
6, 9, 11
0, 7, 12
7, 9, 13
0, 8, 14
8, 9, 15
9, 0, 16
'''
# set all edges
g.set_e_repr(th.zeros((17, D)))
assert _pfc(g.get_e_repr()) == [0.] * 17
# pop edges
assert _pfc(g.pop_e_repr()) == [0.] * 17
assert len(g.get_e_repr()) == 0
g.set_e_repr(th.zeros((17, D)))
# set partial edges (many-many)
u = th.tensor([0, 0, 2, 5, 9])
v = th.tensor([1, 3, 9, 9, 0])
g.set_e_repr(th.ones((5, D)), u, v)
truth = [0.] * 17
truth[0] = truth[4] = truth[3] = truth[9] = truth[16] = 1.
assert _pfc(g.get_e_repr()) == truth
# set partial edges (many-one)
u = th.tensor([3, 4, 6])
v = th.tensor([9])
g.set_e_repr(th.ones((3, D)), u, v)
truth[5] = truth[7] = truth[11] = 1.
assert _pfc(g.get_e_repr()) == truth
# set partial edges (one-many)
u = th.tensor([0])
v = th.tensor([4, 5, 6])
g.set_e_repr(th.ones((3, D)), u, v)
truth[6] = truth[8] = truth[10] = 1.
assert _pfc(g.get_e_repr()) == truth
# get partial edges (many-many)
u = th.tensor([0, 6, 0])
v = th.tensor([6, 9, 7])
assert _pfc(g.get_e_repr(u, v)) == [1., 1., 0.]
# get partial edges (many-one)
u = th.tensor([5, 6, 7])
v = th.tensor([9])
assert _pfc(g.get_e_repr(u, v)) == [1., 1., 0.]
# get partial edges (one-many)
u = th.tensor([0])
v = th.tensor([3, 4, 5])
assert _pfc(g.get_e_repr(u, v)) == [1., 1., 1.]
def test_batch_setter_autograd():
g = generate_graph(grad=True)
h1 = g.get_n_repr()
# partial set
v = th.tensor([1, 2, 8])
hh = Variable(th.zeros((len(v), D)), requires_grad=True)
g.set_n_repr(hh, v)
h2 = g.get_n_repr()
h2.backward(th.ones((10, D)) * 2)
check_eq(h1.grad[:,0], th.tensor([2., 0., 0., 2., 2., 2., 2., 2., 0., 2.]))
check_eq(hh.grad[:,0], th.tensor([2., 2., 2.]))
def test_batch_send():
g = generate_graph()
def _fmsg(hu, edge):
assert hu.shape == (5, D)
return hu
g.register_message_func(_fmsg)
# many-many send
u = th.tensor([0, 0, 0, 0, 0])
v = th.tensor([1, 2, 3, 4, 5])
g.send(u, v)
# one-many send
u = th.tensor([0])
v = th.tensor([1, 2, 3, 4, 5])
g.send(u, v)
# many-one send
u = th.tensor([1, 2, 3, 4, 5])
v = th.tensor([9])
g.send(u, v)
def test_batch_recv():
g = generate_graph()
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
u = th.tensor([0, 0, 0, 4, 5, 6])
v = th.tensor([1, 2, 3, 9, 9, 9])
reduce_msg_shapes.clear()
g.send(u, v)
g.recv(th.unique(v))
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
def test_update_routines():
g = generate_graph()
g.register_message_func(message_func)
g.register_reduce_func(reduce_func)
# send_and_recv
reduce_msg_shapes.clear()
u = th.tensor([0, 0, 0, 4, 5, 6])
v = th.tensor([1, 2, 3, 9, 9, 9])
g.send_and_recv(u, v)
assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
reduce_msg_shapes.clear()
# pull
v = th.tensor([1, 2, 3, 9])
reduce_msg_shapes.clear()
g.pull(v)
assert(reduce_msg_shapes == {(1, 8, D), (3, 1, D)})
reduce_msg_shapes.clear()
# push
v = th.tensor([0, 1, 2, 3])
reduce_msg_shapes.clear()
g.push(v)
assert(reduce_msg_shapes == {(1, 3, D), (8, 1, D)})
reduce_msg_shapes.clear()
# update_all
reduce_msg_shapes.clear()
g.update_all()
assert(reduce_msg_shapes == {(1, 8, D), (9, 1, D)})
reduce_msg_shapes.clear()
if __name__ == '__main__':
test_batch_setter_getter()
test_batch_setter_autograd()
test_batch_send()
test_batch_recv()
test_update_routines()
tests/pytorch/test_batched_graph.py
View file @ 9c135fd5
......@@ -18,8 +18,8 @@ def tree1():
g.add_edge(4, 1)
g.add_edge(1, 0)
g.add_edge(2, 0)
g.set_n_repr(th.Tensor([0, 1, 2, 3, 4]))
g.set_e_repr(th.randn(4, 10))
g.set_n_repr({'h' : th.Tensor([0, 1, 2, 3, 4])})
g.set_e_repr({'h' : th.randn(4, 10)})
return g
def tree2():
......@@ -37,17 +37,17 @@ def tree2():
g.add_edge(0, 4)
g.add_edge(4, 1)
g.add_edge(3, 1)
g.set_n_repr(th.Tensor([0, 1, 2, 3, 4]))
g.set_e_repr(th.randn(4, 10))
g.set_n_repr({'h' : th.Tensor([0, 1, 2, 3, 4])})
g.set_e_repr({'h' : th.randn(4, 10)})
return g
def test_batch_unbatch():
t1 = tree1()
t2 = tree2()
n1 = t1.get_n_repr()
n2 = t2.get_n_repr()
e1 = t1.get_e_repr()
e2 = t2.get_e_repr()
n1 = t1.get_n_repr()['h']
n2 = t2.get_n_repr()['h']
e1 = t1.get_e_repr()['h']
e2 = t2.get_e_repr()['h']
bg = dgl.batch([t1, t2])
assert bg.number_of_nodes() == 10
......@@ -57,10 +57,10 @@ def test_batch_unbatch():
assert bg.batch_num_edges == [4, 4]
tt1, tt2 = dgl.unbatch(bg)
assert th.allclose(t1.get_n_repr(), tt1.get_n_repr())
assert th.allclose(t1.get_e_repr(), tt1.get_e_repr())
assert th.allclose(t2.get_n_repr(), tt2.get_n_repr())
assert th.allclose(t2.get_e_repr(), tt2.get_e_repr())
assert th.allclose(t1.get_n_repr()['h'], tt1.get_n_repr()['h'])
assert th.allclose(t1.get_e_repr()['h'], tt1.get_e_repr()['h'])
assert th.allclose(t2.get_n_repr()['h'], tt2.get_n_repr()['h'])
assert th.allclose(t2.get_e_repr()['h'], tt2.get_e_repr()['h'])
def test_batch_unbatch1():
t1 = tree1()
......@@ -74,20 +74,20 @@ def test_batch_unbatch1():
assert b2.batch_num_edges == [4, 4, 4]
s1, s2, s3 = dgl.unbatch(b2)
assert th.allclose(t2.get_n_repr(), s1.get_n_repr())
assert th.allclose(t2.get_e_repr(), s1.get_e_repr())
assert th.allclose(t1.get_n_repr(), s2.get_n_repr())
assert th.allclose(t1.get_e_repr(), s2.get_e_repr())
assert th.allclose(t2.get_n_repr(), s3.get_n_repr())
assert th.allclose(t2.get_e_repr(), s3.get_e_repr())
assert th.allclose(t2.get_n_repr()['h'], s1.get_n_repr()['h'])
assert th.allclose(t2.get_e_repr()['h'], s1.get_e_repr()['h'])
assert th.allclose(t1.get_n_repr()['h'], s2.get_n_repr()['h'])
assert th.allclose(t1.get_e_repr()['h'], s2.get_e_repr()['h'])
assert th.allclose(t2.get_n_repr()['h'], s3.get_n_repr()['h'])
assert th.allclose(t2.get_e_repr()['h'], s3.get_e_repr()['h'])
def test_batch_sendrecv():
t1 = tree1()
t2 = tree2()
bg = dgl.batch([t1, t2])
bg.register_message_func(lambda src, edge: src)
bg.register_reduce_func(lambda node, msgs: th.sum(msgs, 1))
bg.register_message_func(lambda src, edge: {'m' : src['h']})
bg.register_reduce_func(lambda node, msgs: {'h' : th.sum(msgs['m'], 1)})
u = [3, 4, 2 + 5, 0 + 5]
v = [1, 1, 4 + 5, 4 + 5]
......@@ -95,8 +95,8 @@ def test_batch_sendrecv():
bg.recv(v)
t1, t2 = dgl.unbatch(bg)
assert t1.get_n_repr()[1] == 7
assert t2.get_n_repr()[4] == 2
assert t1.get_n_repr()['h'][1] == 7
assert t2.get_n_repr()['h'][4] == 2
def test_batch_propagate():
......@@ -104,8 +104,8 @@ def test_batch_propagate():
t2 = tree2()
bg = dgl.batch([t1, t2])
bg.register_message_func(lambda src, edge: src)
bg.register_reduce_func(lambda node, msgs: th.sum(msgs, 1))
bg.register_message_func(lambda src, edge: {'m' : src['h']})
bg.register_reduce_func(lambda node, msgs: {'h' : th.sum(msgs['m'], 1)})
# get leaves.
order = []
......@@ -123,45 +123,37 @@ def test_batch_propagate():
bg.propagate(traverser=order)
t1, t2 = dgl.unbatch(bg)
assert t1.get_n_repr()[0] == 9
assert t2.get_n_repr()[1] == 5
assert t1.get_n_repr()['h'][0] == 9
assert t2.get_n_repr()['h'][1] == 5
def test_batched_edge_ordering():
g1 = dgl.DGLGraph()
g1.add_nodes(6)
g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
e1 = th.randn(5, 10)
g1.set_e_repr(e1)
g1.set_e_repr({'h' : e1})
g2 = dgl.DGLGraph()
g2.add_nodes(6)
g2.add_edges([0, 1 ,2 ,5, 4 ,5], [1, 2, 3, 4, 3, 0])
e2 = th.randn(6, 10)
g2.set_e_repr(e2)
g2.set_e_repr({'h' : e2})
g = dgl.batch([g1, g2])
r1 = g.get_e_repr()[g.edge_id(4, 5)]
r2 = g1.get_e_repr()[g1.edge_id(4, 5)]
r1 = g.get_e_repr()['h'][g.edge_id(4, 5)]
r2 = g1.get_e_repr()['h'][g1.edge_id(4, 5)]
assert th.equal(r1, r2)
def test_batch_no_edge():
# FIXME: current impl cannot handle this case!!!
# comment out for now to test CI
return
"""
g1 = dgl.DGLGraph()
g1.add_nodes(6)
g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
e1 = th.randn(5, 10)
g1.set_e_repr(e1)
g2 = dgl.DGLGraph()
g2.add_nodes(6)
g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
e2 = th.randn(6, 10)
g2.set_e_repr(e2)
g3 = dgl.DGLGraph()
g3.add_nodes(1) # no edges
g = dgl.batch([g1, g3, g2]) # should not throw an error
"""
if __name__ == '__main__':
test_batch_unbatch()
......
tests/pytorch/test_frame.py
View file @ 9c135fd5
......@@ -2,14 +2,11 @@ import torch as th
from torch.autograd import Variable
import numpy as np
from dgl.frame import Frame, FrameRef
from dgl.utils import Index
from dgl.utils import Index, toindex
N = 10
D = 5
def check_eq(a, b):
return a.shape == b.shape and np.allclose(a.numpy(), b.numpy())
def check_fail(fn):
try:
fn()
......@@ -27,12 +24,13 @@ def test_create():
data = create_test_data()
f1 = Frame()
for k, v in data.items():
f1.add_column(k, v)
assert f1.schemes == set(data.keys())
f1.update_column(k, v)
print(f1.schemes)
assert f1.keys() == set(data.keys())
assert f1.num_columns == 3
assert f1.num_rows == N
f2 = Frame(data)
assert f2.schemes == set(data.keys())
assert f2.keys() == set(data.keys())
assert f2.num_columns == 3
assert f2.num_rows == N
f1.clear()
......@@ -45,9 +43,9 @@ def test_column1():
f = Frame(data)
assert f.num_rows == N
assert len(f) == 3
assert check_eq(f['a1'], data['a1'])
assert th.allclose(f['a1'].data, data['a1'].data)
f['a1'] = data['a2']
assert check_eq(f['a2'], data['a2'])
assert th.allclose(f['a2'].data, data['a2'].data)
# add a different length column should fail
def failed_add_col():
f['a4'] = th.zeros([N+1, D])
......@@ -70,16 +68,15 @@ def test_column2():
f = FrameRef(data, [3, 4, 5, 6, 7])
assert f.num_rows == 5
assert len(f) == 3
assert check_eq(f['a1'], data['a1'][3:8])
assert th.allclose(f['a1'], data['a1'].data[3:8])
# set column should reflect on the referenced data
f['a1'] = th.zeros([5, D])
assert check_eq(data['a1'][3:8], th.zeros([5, D]))
# add new column should be padded with zero
assert th.allclose(data['a1'].data[3:8], th.zeros([5, D]))
# add new partial column should fail with error initializer
f.set_initializer(lambda shape, dtype : assert_(False))
def failed_add_col():
f['a4'] = th.ones([5, D])
assert len(data) == 4
assert check_eq(data['a4'][0:3], th.zeros([3, D]))
assert check_eq(data['a4'][3:8], th.ones([5, D]))
assert check_eq(data['a4'][8:10], th.zeros([2, D]))
assert check_fail(failed_add_col)
def test_append1():
# test append API on Frame
......@@ -91,9 +88,14 @@ def test_append1():
f1.append(f2)
assert f1.num_rows == 2 * N
c1 = f1['a1']
assert c1.shape == (2 * N, D)
assert c1.data.shape == (2 * N, D)
truth = th.cat([data['a1'], data['a1']])
assert check_eq(truth, c1)
assert th.allclose(truth, c1.data)
# append dict of different length columns should fail
f3 = {'a1' : th.zeros((3, D)), 'a2' : th.zeros((3, D)), 'a3' : th.zeros((2, D))}
def failed_append():
f1.append(f3)
assert check_fail(failed_append)
def test_append2():
# test append on FrameRef
......@@ -113,7 +115,7 @@ def test_append2():
assert not f.is_span_whole_column()
assert f.num_rows == 3 * N
new_idx = list(range(N)) + list(range(2*N, 4*N))
assert check_eq(f.index().tousertensor(), th.tensor(new_idx))
assert th.all(f.index().tousertensor() == th.tensor(new_idx, dtype=th.int64))
assert data.num_rows == 4 * N
def test_row1():
......@@ -127,13 +129,13 @@ def test_row1():
rows = f[rowid]
for k, v in rows.items():
assert v.shape == (len(rowid), D)
assert check_eq(v, data[k][rowid])
assert th.allclose(v, data[k][rowid])
# test duplicate keys
rowid = Index(th.tensor([8, 2, 2, 1]))
rows = f[rowid]
for k, v in rows.items():
assert v.shape == (len(rowid), D)
assert check_eq(v, data[k][rowid])
assert th.allclose(v, data[k][rowid])
# setter
rowid = Index(th.tensor([0, 2, 4]))
......@@ -143,12 +145,14 @@ def test_row1():
}
f[rowid] = vals
for k, v in f[rowid].items():
assert check_eq(v, th.zeros((len(rowid), D)))
assert th.allclose(v, th.zeros((len(rowid), D)))
# setting rows with new column should automatically add a new column
# setting rows with new column should raise error with error initializer
f.set_initializer(lambda shape, dtype : assert_(False))
def failed_update_rows():
vals['a4'] = th.ones((len(rowid), D))
f[rowid] = vals
assert len(f) == 4
assert check_fail(failed_update_rows)
def test_row2():
# test row getter/setter autograd compatibility
......@@ -161,13 +165,13 @@ def test_row2():
rowid = Index(th.tensor([0, 2]))
rows = f[rowid]
rows['a1'].backward(th.ones((len(rowid), D)))
assert check_eq(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
assert th.allclose(c1.grad[:,0], th.tensor([1., 0., 1., 0., 0., 0., 0., 0., 0., 0.]))
c1.grad.data.zero_()
# test duplicate keys
rowid = Index(th.tensor([8, 2, 2, 1]))
rows = f[rowid]
rows['a1'].backward(th.ones((len(rowid), D)))
assert check_eq(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 2., 0., 0., 0., 0., 0., 1., 0.]))
c1.grad.data.zero_()
# setter
......@@ -180,8 +184,8 @@ def test_row2():
f[rowid] = vals
c11 = f['a1']
c11.backward(th.ones((N, D)))
assert check_eq(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
assert check_eq(vals['a1'].grad, th.ones((len(rowid), D)))
assert th.allclose(c1.grad[:,0], th.tensor([0., 1., 0., 1., 0., 1., 1., 1., 1., 1.]))
assert th.allclose(vals['a1'].grad, th.ones((len(rowid), D)))
assert vals['a2'].grad is None
def test_row3():
......@@ -201,8 +205,9 @@ def test_row3():
newidx = list(range(N))
newidx.pop(2)
newidx.pop(2)
newidx = toindex(newidx)
for k, v in f.items():
assert check_eq(v, data[k][th.tensor(newidx)])
assert th.allclose(v, data[k][newidx])
def test_sharing():
data = Frame(create_test_data())
......@@ -210,10 +215,10 @@ def test_sharing():
f2 = FrameRef(data, index=[2, 3, 4, 5, 6])
# test read
for k, v in f1.items():
assert check_eq(data[k][0:4], v)
assert th.allclose(data[k].data[0:4], v)
for k, v in f2.items():
assert check_eq(data[k][2:7], v)
f2_a1 = f2['a1']
assert th.allclose(data[k].data[2:7], v)
f2_a1 = f2['a1'].data
# test write
# update own ref should not been seen by the other.
f1[Index(th.tensor([0, 1]))] = {
......@@ -221,7 +226,7 @@ def test_sharing():
'a2' : th.zeros([2, D]),
'a3' : th.zeros([2, D]),
}
assert check_eq(f2['a1'], f2_a1)
assert th.allclose(f2['a1'], f2_a1)
# update shared space should been seen by the other.
f1[Index(th.tensor([2, 3]))] = {
'a1' : th.ones([2, D]),
......@@ -229,7 +234,7 @@ def test_sharing():
'a3' : th.ones([2, D]),
}
f2_a1[0:2] = th.ones([2, D])
assert check_eq(f2['a1'], f2_a1)
assert th.allclose(f2['a1'], f2_a1)
if __name__ == '__main__':
test_create()
......
tests/pytorch/test_function.py
View file @ 9c135fd5
import torch as th
import dgl
import dgl.function as fn
from dgl.graph import __REPR__
def generate_graph():
g = dgl.DGLGraph()
......@@ -37,18 +36,9 @@ def generate_graph1():
g.set_e_repr(h)
return g
def reducer_msg(node, msgs):
return th.sum(msgs['m'], 1)
def reducer_out(node, msgs):
return {'h' : th.sum(msgs, 1)}
def reducer_both(node, msgs):
return {'h' : th.sum(msgs['m'], 1)}
def reducer_none(node, msgs):
return th.sum(msgs, 1)
def test_copy_src():
# copy_src with both fields
g = generate_graph()
......@@ -58,30 +48,6 @@ def test_copy_src():
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy_src with only src field; the out field should use anonymous repr
g = generate_graph()
g.register_message_func(fn.copy_src(src='h'))
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy_src with no src field; should use anonymous repr
g = generate_graph1()
g.register_message_func(fn.copy_src(out='m'))
g.register_reduce_func(reducer_both)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy src with no fields;
g = generate_graph1()
g.register_message_func(fn.copy_src())
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
def test_copy_edge():
# copy_edge with both fields
g = generate_graph()
......@@ -91,30 +57,6 @@ def test_copy_edge():
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy_edge with only edge field; the out field should use anonymous repr
g = generate_graph()
g.register_message_func(fn.copy_edge(edge='h'))
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy_edge with no edge field; should use anonymous repr
g = generate_graph1()
g.register_message_func(fn.copy_edge(out='m'))
g.register_reduce_func(reducer_both)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
# copy edge with no fields;
g = generate_graph1()
g.register_message_func(fn.copy_edge())
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([10., 1., 1., 1., 1., 1., 1., 1., 1., 44.]))
def test_src_mul_edge():
# src_mul_edge with all fields
g = generate_graph()
......@@ -124,34 +66,6 @@ def test_src_mul_edge():
assert th.allclose(g.get_n_repr()['h'],
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
g = generate_graph()
g.register_message_func(fn.src_mul_edge(src='h', edge='h'))
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
g = generate_graph1()
g.register_message_func(fn.src_mul_edge(out='m'))
g.register_reduce_func(reducer_both)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
g = generate_graph1()
g.register_message_func(fn.src_mul_edge())
g.register_reduce_func(reducer_out)
g.update_all()
assert th.allclose(g.get_n_repr()['h'],
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
g = generate_graph1()
g.register_message_func(fn.src_mul_edge())
g.register_reduce_func(reducer_none)
g.update_all()
assert th.allclose(g.get_n_repr(),
th.tensor([100., 1., 1., 1., 1., 1., 1., 1., 1., 284.]))
if __name__ == '__main__':
test_copy_src()
test_copy_edge()
......
tests/pytorch/test_line_graph.py
View file @ 9c135fd5
......@@ -5,35 +5,31 @@ import dgl
D = 5
def check_eq(a, b):
return a.shape == b.shape and np.allclose(a.numpy(), b.numpy())
def test_line_graph():
N = 5
G = dgl.DGLGraph(nx.star_graph(N))
G.set_e_repr(th.randn((2 * N, D)))
G.set_e_repr({'h' : th.randn((2 * N, D))})
n_edges = G.number_of_edges()
L = G.line_graph(shared=True)
assert L.number_of_nodes() == 2 * N
L.set_n_repr(th.randn((2 * N, D)))
L.set_n_repr({'h' : th.randn((2 * N, D))})
# update node features on line graph should reflect to edge features on
# original graph.
u = [0, 0, 2, 3]
v = [1, 2, 0, 0]
eid = G.edge_ids(u, v)
L.set_n_repr(th.zeros((4, D)), eid)
assert check_eq(G.get_e_repr(u, v), th.zeros((4, D)))
L.set_n_repr({'h' : th.zeros((4, D))}, eid)
assert th.allclose(G.get_e_repr(u, v)['h'], th.zeros((4, D)))
# adding a new node feature on line graph should also reflect to a new
# edge feature on original graph
data = th.randn(n_edges, D)
L.set_n_repr({'w': data})
assert check_eq(G.get_e_repr()['w'], data)
assert th.allclose(G.get_e_repr()['w'], data)
def test_no_backtracking():
N = 5
G = dgl.DGLGraph(nx.star_graph(N))
G.set_e_repr(th.randn((2 * N, D)))
L = G.line_graph(backtracking=False)
assert L.number_of_nodes() == 2 * N
for i in range(1, N):
......
tests/pytorch/test_specialization.py
View file @ 9c135fd5
......@@ -22,23 +22,23 @@ def generate_graph():
def test_update_all():
def _test(fld):
def message_func(hu, edge):
return hu[fld]
return {'m' : hu[fld]}
def message_func_edge(hu, edge):
if len(hu[fld].shape) == 1:
return hu[fld] * edge['e1']
return {'m' : hu[fld] * edge['e1']}
else:
return hu[fld] * edge['e2']
return {'m' : hu[fld] * edge['e2']}
def reduce_func(hv, msgs):
return {fld : th.sum(msgs, 1)}
return {fld : th.sum(msgs['m'], 1)}
def apply_func(hu):
return {fld : 2 * hu[fld]}
g = generate_graph()
# update all
v1 = g.get_n_repr()[fld]
g.update_all(fn.copy_src(src=fld), fn.sum(out=fld), apply_func)
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out=fld), apply_func)
v2 = g.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.update_all(message_func, reduce_func, apply_func)
......@@ -46,12 +46,12 @@ def test_update_all():
assert th.allclose(v2, v3)
# update all with edge weights
v1 = g.get_n_repr()[fld]
g.update_all(fn.src_mul_edge(src=fld, edge='e1'),
fn.sum(out=fld), apply_func)
g.update_all(fn.src_mul_edge(src=fld, edge='e1', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.update_all(fn.src_mul_edge(src=fld, edge='e2'),
fn.sum(out=fld), apply_func)
g.update_all(fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.update_all(message_func_edge, reduce_func, apply_func)
......@@ -68,42 +68,40 @@ def test_send_and_recv():
v = th.tensor([1, 2, 3, 9, 9, 0])
def _test(fld):
def message_func(hu, edge):
return hu[fld]
return {'m' : hu[fld]}
def message_func_edge(hu, edge):
if len(hu[fld].shape) == 1:
return hu[fld] * edge['e1']
return {'m' : hu[fld] * edge['e1']}
else:
return hu[fld] * edge['e2']
return {'m' : hu[fld] * edge['e2']}
def reduce_func(hv, msgs):
return {fld : th.sum(msgs, 1)}
return {fld : th.sum(msgs['m'], 1)}
def apply_func(hu):
return {fld : 2 * hu[fld]}
g = generate_graph()
# send and recv
v1 = g.get_n_repr()[fld]
g.send_and_recv(u, v, fn.copy_src(src=fld),
fn.sum(out=fld), apply_func)
g.send_and_recv(u, v, fn.copy_src(src=fld, out='m'),
fn.sum(msg='m', out=fld), apply_func)
v2 = g.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.send_and_recv(u, v, message_func,
reduce_func, apply_func)
g.send_and_recv(u, v, message_func, reduce_func, apply_func)
v3 = g.get_n_repr()[fld]
assert th.allclose(v2, v3)
# send and recv with edge weights
v1 = g.get_n_repr()[fld]
g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e1'),
fn.sum(out=fld), apply_func)
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.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e2'),
fn.sum(out=fld), apply_func)
g.send_and_recv(u, v, fn.src_mul_edge(src=fld, edge='e2', out='m'),
fn.sum(msg='m', out=fld), apply_func)
v3 = g.get_n_repr()[fld]
g.set_n_repr({fld : v1})
g.send_and_recv(u, v, message_func_edge,
reduce_func, apply_func)
g.send_and_recv(u, v, message_func_edge, reduce_func, apply_func)
v4 = g.get_n_repr()[fld]
assert th.allclose(v2, v3)
assert th.allclose(v3, v4)
......@@ -123,22 +121,23 @@ def test_update_all_multi_fn():
return {'v2': th.sum(msgs['m2'], 1)}
g = generate_graph()
g.set_n_repr({'v1' : th.zeros((10,)), 'v2' : th.zeros((10,))})
fld = 'f2'
# update all, mix of builtin and UDF
g.update_all([fn.copy_src(src=fld, out='m1'), message_func],
[fn.sum(msgs='m1', out='v1'), reduce_func],
[fn.sum(msg='m1', out='v1'), reduce_func],
None)
v1 = g.get_n_repr()['v1']
v2 = g.get_n_repr()['v2']
assert th.allclose(v1, v2)
# run builtin with single message and reduce
g.update_all(fn.copy_src(src=fld), fn.sum(out='v1'), None)
g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'), None)
v1 = g.get_n_repr()['v1']
assert th.allclose(v1, v2)
# 1 message, 2 reduces, using anonymous repr
g.update_all(fn.copy_src(src=fld), [fn.sum(out='v2'), fn.sum(out='v3')], None)
# 1 message, 2 reduces
g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')], None)
v2 = g.get_n_repr()['v2']
v3 = g.get_n_repr()['v3']
assert th.allclose(v1, v2)
......@@ -146,7 +145,7 @@ def test_update_all_multi_fn():
# update all with edge weights, 2 message, 3 reduces
g.update_all([fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
[fn.sum(msgs='m1', out='v1'), fn.sum(msgs='m2', out='v2'), fn.sum(msgs='m1', out='v3')],
[fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
None)
v1 = g.get_n_repr()['v1']
v2 = g.get_n_repr()['v2']
......@@ -173,25 +172,30 @@ def test_send_and_recv_multi_fn():
return {'v2' : th.sum(msgs['m2'], 1)}
g = generate_graph()
g.set_n_repr({'v1' : th.zeros((10, D)), 'v2' : th.zeros((10, D)),
'v3' : th.zeros((10, D))})
fld = 'f2'
# send and recv, mix of builtin and UDF
g.send_and_recv(u, v,
[fn.copy_src(src=fld, out='m1'), message_func],
[fn.sum(msgs='m1', out='v1'), reduce_func],
[fn.sum(msg='m1', out='v1'), reduce_func],
None)
v1 = g.get_n_repr()['v1']
v2 = g.get_n_repr()['v2']
assert th.allclose(v1, v2)
# run builtin with single message and reduce
g.send_and_recv(u, v, fn.copy_src(src=fld), fn.sum(out='v1'),
g.send_and_recv(u, v, fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'),
None)
v1 = g.get_n_repr()['v1']
assert th.allclose(v1, v2)
# 1 message, 2 reduces, using anonymous repr
g.send_and_recv(u, v, fn.copy_src(src=fld), [fn.sum(out='v2'), fn.sum(out='v3')], None)
# 1 message, 2 reduces
g.send_and_recv(u, v,
fn.copy_src(src=fld, out='m'),
[fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')],
None)
v2 = g.get_n_repr()['v2']
v3 = g.get_n_repr()['v3']
assert th.allclose(v1, v2)
......@@ -200,7 +204,7 @@ def test_send_and_recv_multi_fn():
# send and recv with edge weights, 2 message, 3 reduces
g.send_and_recv(u, v,
[fn.src_mul_edge(src=fld, edge='e1', out='m1'), fn.src_mul_edge(src=fld, edge='e2', out='m2')],
[fn.sum(msgs='m1', out='v1'), fn.sum(msgs='m2', out='v2'), fn.sum(msgs='m1', out='v3')],
[fn.sum(msg='m1', out='v1'), fn.sum(msg='m2', out='v2'), fn.sum(msg='m1', out='v3')],
None)
v1 = g.get_n_repr()['v1']
v2 = g.get_n_repr()['v2']
......
tests/scripts/task_lint.sh 0 → 100755
View file @ 9c135fd5
#!/bin/sh
# cpplint
echo 'Checking code style of C++ codes...'
python3 third_party/dmlc-core/scripts/lint.py dgl cpp include src
tests/scripts/test_examples.sh 0 → 100755
View file @ 9c135fd5
#!/bin/bash
GCN_EXAMPLE_DIR="../../examples/pytorch/gcn"
function fail {
echo FAIL: $@
exit -1
}
function usage {
echo "Usage: $0 [CPU|GPU]"
}
# check arguments
if [ $# -ne 1 ]; then
usage
fail "Error: must specify device"
fi
if [ "$1" == "CPU" ]; then
dev=-1
elif [ "$1" == "GPU" ]; then
export CUDA_VISIBLE_DEVICES=0
dev=0
else
usage
fail "Unknown device $1"
fi
pushd $GCN_EXAMPLE_DIR> /dev/null
# test CPU
python3 gcn.py --dataset cora --gpu $dev || fail "run gcn.py on $1"
python3 gcn_spmv.py --dataset cora --gpu $dev || fail "run gcn_spmv.py on $1"
popd > /dev/null
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment