Skip to content

GitLab

Unverified Commit 89a4cc4d authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] Black auto fix. (#4694) 


Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent 303b150f
Hide whitespace changes
Inline Side-by-side
Showing with 754 additions and 433 deletions
tests/mxnet/test_geometry.py
View file @ 89a4cc4d
import backend as F
import mxnet as mx import mxnet as mx
import numpy as np
from dgl.geometry import farthest_point_sampler from dgl.geometry import farthest_point_sampler
import backend as F
import numpy as np
def test_fps(): def test_fps():
N = 1000 N = 1000
batch_size = 5 batch_size = 5
sample_points = 10 sample_points = 10
x = mx.nd.array(np.random.uniform(size=(batch_size, int(N/batch_size), 3))) x = mx.nd.array(
np.random.uniform(size=(batch_size, int(N / batch_size), 3))
)
ctx = F.ctx() ctx = F.ctx()
if F.gpu_ctx(): if F.gpu_ctx():
x = x.as_in_context(ctx) x = x.as_in_context(ctx)
...@@ -17,5 +20,6 @@ def test_fps(): ...@@ -17,5 +20,6 @@ def test_fps():
assert res.shape[1] == sample_points assert res.shape[1] == sample_points
assert res.sum() > 0 assert res.sum() > 0
if __name__ == '__main__':
if __name__ == "__main__":
test_fps() test_fps()
tests/mxnet/test_nn.py
View file @ 89a4cc4d
import backend as F
import mxnet as mx import mxnet as mx
import networkx as nx import networkx as nx
import numpy as np import numpy as np
import scipy as sp
import pytest import pytest
import scipy as sp
from mxnet import autograd, gluon, nd
from test_utils import parametrize_idtype
from test_utils.graph_cases import (
get_cases,
random_bipartite,
random_dglgraph,
random_graph,
)
import dgl import dgl
import dgl.nn.mxnet as nn
import dgl.function as fn import dgl.function as fn
import backend as F import dgl.nn.mxnet as nn
from test_utils.graph_cases import get_cases, random_graph, random_bipartite, random_dglgraph
from test_utils import parametrize_idtype
from mxnet import autograd, gluon, nd
def check_close(a, b): def check_close(a, b):
assert np.allclose(a.asnumpy(), b.asnumpy(), rtol=1e-4, atol=1e-4) assert np.allclose(a.asnumpy(), b.asnumpy(), rtol=1e-4, atol=1e-4)
def _AXWb(A, X, W, b): def _AXWb(A, X, W, b):
X = mx.nd.dot(X, W.data(X.context)) X = mx.nd.dot(X, W.data(X.context))
Y = mx.nd.dot(A, X.reshape(X.shape[0], -1)).reshape(X.shape) Y = mx.nd.dot(A, X.reshape(X.shape[0], -1)).reshape(X.shape)
return Y + b.data(X.context) return Y + b.data(X.context)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("out_dim", [1, 2])
def test_graph_conv(idtype, out_dim): def test_graph_conv(idtype, out_dim):
g = dgl.from_networkx(nx.path_graph(3)) g = dgl.from_networkx(nx.path_graph(3))
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
adj = g.adjacency_matrix(transpose=True, ctx=ctx) adj = g.adjacency_matrix(transpose=True, ctx=ctx)
conv = nn.GraphConv(5, out_dim, norm='none', bias=True) conv = nn.GraphConv(5, out_dim, norm="none", bias=True)
conv.initialize(ctx=ctx) conv.initialize(ctx=ctx)
# test#1: basic # test#1: basic
h0 = F.ones((3, 5)) h0 = F.ones((3, 5))
...@@ -77,14 +86,18 @@ def test_graph_conv(idtype, out_dim): ...@@ -77,14 +86,18 @@ def test_graph_conv(idtype, out_dim):
assert len(g.ndata) == 1 assert len(g.ndata) == 1
assert len(g.edata) == 0 assert len(g.edata) == 0
assert "h" in g.ndata assert "h" in g.ndata
check_close(g.ndata['h'], 2 * F.ones((3, 1))) check_close(g.ndata["h"], 2 * F.ones((3, 1)))
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('norm', ['none', 'both', 'right', 'left']) "g",
@pytest.mark.parametrize('weight', [True, False]) get_cases(["homo", "block-bipartite"], exclude=["zero-degree", "dglgraph"]),
@pytest.mark.parametrize('bias', [False]) )
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("norm", ["none", "both", "right", "left"])
@pytest.mark.parametrize("weight", [True, False])
@pytest.mark.parametrize("bias", [False])
@pytest.mark.parametrize("out_dim", [1, 2])
def test_graph_conv2(idtype, g, norm, weight, bias, out_dim): def test_graph_conv2(idtype, g, norm, weight, bias, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
conv = nn.GraphConv(5, out_dim, norm=norm, weight=weight, bias=bias) conv = nn.GraphConv(5, out_dim, norm=norm, weight=weight, bias=bias)
...@@ -99,12 +112,15 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim): ...@@ -99,12 +112,15 @@ def test_graph_conv2(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, h, ext_w) h_out = conv(g, h, ext_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree', 'dglgraph'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('norm', ['none', 'both', 'right']) "g", get_cases(["bipartite"], exclude=["zero-degree", "dglgraph"])
@pytest.mark.parametrize('weight', [True, False]) )
@pytest.mark.parametrize('bias', [False]) @pytest.mark.parametrize("norm", ["none", "both", "right"])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("weight", [True, False])
@pytest.mark.parametrize("bias", [False])
@pytest.mark.parametrize("out_dim", [1, 2])
def test_graph_conv2_bi(idtype, g, norm, weight, bias, out_dim): def test_graph_conv2_bi(idtype, g, norm, weight, bias, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
conv = nn.GraphConv(5, out_dim, norm=norm, weight=weight, bias=bias) conv = nn.GraphConv(5, out_dim, norm=norm, weight=weight, bias=bias)
...@@ -120,6 +136,7 @@ def test_graph_conv2_bi(idtype, g, norm, weight, bias, out_dim): ...@@ -120,6 +136,7 @@ def test_graph_conv2_bi(idtype, g, norm, weight, bias, out_dim):
h_out = conv(g, (h, h_dst), ext_w) h_out = conv(g, (h, h_dst), ext_w)
assert h_out.shape == (ndst, out_dim) assert h_out.shape == (ndst, out_dim)
def _S2AXWb(A, N, X, W, b): def _S2AXWb(A, N, X, W, b):
X1 = X * N X1 = X * N
X1 = mx.nd.dot(A, X1.reshape(X1.shape[0], -1)) X1 = mx.nd.dot(A, X1.reshape(X1.shape[0], -1))
...@@ -132,12 +149,13 @@ def _S2AXWb(A, N, X, W, b): ...@@ -132,12 +149,13 @@ def _S2AXWb(A, N, X, W, b):
return Y + b return Y + b
@pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize("out_dim", [1, 2])
def test_tagconv(out_dim): def test_tagconv(out_dim):
g = dgl.from_networkx(nx.path_graph(3)).to(F.ctx()) g = dgl.from_networkx(nx.path_graph(3)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
adj = g.adjacency_matrix(transpose=True, ctx=ctx) adj = g.adjacency_matrix(transpose=True, ctx=ctx)
norm = mx.nd.power(g.in_degrees().astype('float32'), -0.5) norm = mx.nd.power(g.in_degrees().astype("float32"), -0.5)
conv = nn.TAGConv(5, out_dim, bias=True) conv = nn.TAGConv(5, out_dim, bias=True)
conv.initialize(ctx=ctx) conv.initialize(ctx=ctx)
...@@ -151,7 +169,9 @@ def test_tagconv(out_dim): ...@@ -151,7 +169,9 @@ def test_tagconv(out_dim):
shp = norm.shape + (1,) * (h0.ndim - 1) shp = norm.shape + (1,) * (h0.ndim - 1)
norm = norm.reshape(shp).as_in_context(h0.context) norm = norm.reshape(shp).as_in_context(h0.context)
assert F.allclose(h1, _S2AXWb(adj, norm, h0, conv.lin.data(ctx), conv.h_bias.data(ctx))) assert F.allclose(
h1, _S2AXWb(adj, norm, h0, conv.lin.data(ctx), conv.h_bias.data(ctx))
)
conv = nn.TAGConv(5, out_dim) conv = nn.TAGConv(5, out_dim)
conv.initialize(ctx=ctx) conv.initialize(ctx=ctx)
...@@ -161,14 +181,17 @@ def test_tagconv(out_dim): ...@@ -161,14 +181,17 @@ def test_tagconv(out_dim):
h1 = conv(g, h0) h1 = conv(g, h0)
assert h1.shape[-1] == out_dim assert h1.shape[-1] == out_dim
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('out_dim', [1, 20]) "g", get_cases(["homo", "block-bipartite"], exclude=["zero-degree"])
@pytest.mark.parametrize('num_heads', [1, 5]) )
@pytest.mark.parametrize("out_dim", [1, 20])
@pytest.mark.parametrize("num_heads", [1, 5])
def test_gat_conv(g, idtype, out_dim, num_heads): def test_gat_conv(g, idtype, out_dim, num_heads):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
gat = nn.GATConv(10, out_dim, num_heads) # n_heads = 5 gat = nn.GATConv(10, out_dim, num_heads) # n_heads = 5
gat.initialize(ctx=ctx) gat.initialize(ctx=ctx)
print(gat) print(gat)
feat = F.randn((g.number_of_src_nodes(), 10)) feat = F.randn((g.number_of_src_nodes(), 10))
...@@ -182,25 +205,30 @@ def test_gat_conv(g, idtype, out_dim, num_heads): ...@@ -182,25 +205,30 @@ def test_gat_conv(g, idtype, out_dim, num_heads):
gat.initialize(ctx=ctx) gat.initialize(ctx=ctx)
h = gat(g, feat) h = gat(g, feat)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize("g", get_cases(["bipartite"], exclude=["zero-degree"]))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("out_dim", [1, 2])
@pytest.mark.parametrize('num_heads', [1, 4]) @pytest.mark.parametrize("num_heads", [1, 4])
def test_gat_conv_bi(g, idtype, out_dim, num_heads): def test_gat_conv_bi(g, idtype, out_dim, num_heads):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
gat = nn.GATConv(5, out_dim, num_heads) gat = nn.GATConv(5, out_dim, num_heads)
gat.initialize(ctx=ctx) gat.initialize(ctx=ctx)
feat = (F.randn((g.number_of_src_nodes(), 5)), F.randn((g.number_of_dst_nodes(), 5))) feat = (
F.randn((g.number_of_src_nodes(), 5)),
F.randn((g.number_of_dst_nodes(), 5)),
)
h = gat(g, feat) h = gat(g, feat)
assert h.shape == (g.number_of_dst_nodes(), num_heads, out_dim) assert h.shape == (g.number_of_dst_nodes(), num_heads, out_dim)
_, a = gat(g, feat, True) _, a = gat(g, feat, True)
assert a.shape == (g.number_of_edges(), num_heads, 1) assert a.shape == (g.number_of_edges(), num_heads, 1)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize("g", get_cases(["homo", "block-bipartite"]))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize("aggre_type", ["mean", "pool", "gcn"])
@pytest.mark.parametrize('out_dim', [1, 10]) @pytest.mark.parametrize("out_dim", [1, 10])
def test_sage_conv(idtype, g, aggre_type, out_dim): def test_sage_conv(idtype, g, aggre_type, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -210,36 +238,41 @@ def test_sage_conv(idtype, g, aggre_type, out_dim): ...@@ -210,36 +238,41 @@ def test_sage_conv(idtype, g, aggre_type, out_dim):
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize("g", get_cases(["bipartite"]))
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize("aggre_type", ["mean", "pool", "gcn"])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("out_dim", [1, 2])
def test_sage_conv_bi(idtype, g, aggre_type, out_dim): def test_sage_conv_bi(idtype, g, aggre_type, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
dst_dim = 5 if aggre_type != 'gcn' else 10 dst_dim = 5 if aggre_type != "gcn" else 10
sage = nn.SAGEConv((10, dst_dim), out_dim, aggre_type) sage = nn.SAGEConv((10, dst_dim), out_dim, aggre_type)
feat = (F.randn((g.number_of_src_nodes(), 10)), F.randn((g.number_of_dst_nodes(), dst_dim))) feat = (
F.randn((g.number_of_src_nodes(), 10)),
F.randn((g.number_of_dst_nodes(), dst_dim)),
)
sage.initialize(ctx=ctx) sage.initialize(ctx=ctx)
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == g.number_of_dst_nodes() assert h.shape[0] == g.number_of_dst_nodes()
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('aggre_type', ['mean', 'pool', 'gcn']) @pytest.mark.parametrize("aggre_type", ["mean", "pool", "gcn"])
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("out_dim", [1, 2])
def test_sage_conv_bi2(idtype, aggre_type, out_dim): def test_sage_conv_bi2(idtype, aggre_type, out_dim):
# Test the case for graphs without edges # Test the case for graphs without edges
g = dgl.heterograph({('_U', '_E', '_V'): ([], [])}, {'_U': 5, '_V': 3}) g = dgl.heterograph({("_U", "_E", "_V"): ([], [])}, {"_U": 5, "_V": 3})
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
sage = nn.SAGEConv((3, 3), out_dim, 'gcn') sage = nn.SAGEConv((3, 3), out_dim, "gcn")
feat = (F.randn((5, 3)), F.randn((3, 3))) feat = (F.randn((5, 3)), F.randn((3, 3)))
sage.initialize(ctx=ctx) sage.initialize(ctx=ctx)
h = sage(g, feat) h = sage(g, feat)
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == 3 assert h.shape[0] == 3
for aggre_type in ['mean', 'pool']: for aggre_type in ["mean", "pool"]:
sage = nn.SAGEConv((3, 1), out_dim, aggre_type) sage = nn.SAGEConv((3, 1), out_dim, aggre_type)
feat = (F.randn((5, 3)), F.randn((3, 1))) feat = (F.randn((5, 3)), F.randn((3, 1)))
sage.initialize(ctx=ctx) sage.initialize(ctx=ctx)
...@@ -247,11 +280,12 @@ def test_sage_conv_bi2(idtype, aggre_type, out_dim): ...@@ -247,11 +280,12 @@ def test_sage_conv_bi2(idtype, aggre_type, out_dim):
assert h.shape[-1] == out_dim assert h.shape[-1] == out_dim
assert h.shape[0] == 3 assert h.shape[0] == 3
def test_gg_conv(): def test_gg_conv():
g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx()) g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
gg_conv = nn.GatedGraphConv(10, 20, 3, 4) # n_step = 3, n_etypes = 4 gg_conv = nn.GatedGraphConv(10, 20, 3, 4) # n_step = 3, n_etypes = 4
gg_conv.initialize(ctx=ctx) gg_conv.initialize(ctx=ctx)
print(gg_conv) print(gg_conv)
...@@ -261,12 +295,13 @@ def test_gg_conv(): ...@@ -261,12 +295,13 @@ def test_gg_conv():
h1 = gg_conv(g, h0, etypes) h1 = gg_conv(g, h0, etypes)
assert h1.shape == (20, 20) assert h1.shape == (20, 20)
@pytest.mark.parametrize('out_dim', [1, 20])
@pytest.mark.parametrize("out_dim", [1, 20])
def test_cheb_conv(out_dim): def test_cheb_conv(out_dim):
g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx()) g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
cheb = nn.ChebConv(10, out_dim, 3) # k = 3 cheb = nn.ChebConv(10, out_dim, 3) # k = 3
cheb.initialize(ctx=ctx) cheb.initialize(ctx=ctx)
print(cheb) print(cheb)
...@@ -275,8 +310,11 @@ def test_cheb_conv(out_dim): ...@@ -275,8 +310,11 @@ def test_cheb_conv(out_dim):
h1 = cheb(g, h0) h1 = cheb(g, h0)
assert h1.shape == (20, out_dim) assert h1.shape == (20, out_dim)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize(
"g", get_cases(["homo", "block-bipartite"], exclude=["zero-degree"])
)
def test_agnn_conv(g, idtype): def test_agnn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -287,18 +325,23 @@ def test_agnn_conv(g, idtype): ...@@ -287,18 +325,23 @@ def test_agnn_conv(g, idtype):
h = agnn_conv(g, feat) h = agnn_conv(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 10) assert h.shape == (g.number_of_dst_nodes(), 10)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize("g", get_cases(["bipartite"], exclude=["zero-degree"]))
def test_agnn_conv_bi(g, idtype): def test_agnn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
agnn_conv = nn.AGNNConv(0.1, True) agnn_conv = nn.AGNNConv(0.1, True)
agnn_conv.initialize(ctx=ctx) agnn_conv.initialize(ctx=ctx)
print(agnn_conv) print(agnn_conv)
feat = (F.randn((g.number_of_src_nodes(), 5)), F.randn((g.number_of_dst_nodes(), 5))) feat = (
F.randn((g.number_of_src_nodes(), 5)),
F.randn((g.number_of_dst_nodes(), 5)),
)
h = agnn_conv(g, feat) h = agnn_conv(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 5) assert h.shape == (g.number_of_dst_nodes(), 5)
def test_appnp_conv(): def test_appnp_conv():
g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx()) g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -312,69 +355,70 @@ def test_appnp_conv(): ...@@ -312,69 +355,70 @@ def test_appnp_conv():
h1 = appnp_conv(g, h0) h1 = appnp_conv(g, h0)
assert h1.shape == (20, 10) assert h1.shape == (20, 10)
@pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize("out_dim", [1, 2])
def test_dense_cheb_conv(out_dim): def test_dense_cheb_conv(out_dim):
for k in range(1, 4): for k in range(1, 4):
ctx = F.ctx() ctx = F.ctx()
g = dgl.from_scipy(sp.sparse.random(100, 100, density=0.3)).to(F.ctx()) g = dgl.from_scipy(sp.sparse.random(100, 100, density=0.3)).to(F.ctx())
adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype('default') adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype("default")
cheb = nn.ChebConv(5, out_dim, k) cheb = nn.ChebConv(5, out_dim, k)
dense_cheb = nn.DenseChebConv(5, out_dim, k) dense_cheb = nn.DenseChebConv(5, out_dim, k)
cheb.initialize(ctx=ctx) cheb.initialize(ctx=ctx)
dense_cheb.initialize(ctx=ctx) dense_cheb.initialize(ctx=ctx)
for i in range(len(cheb.fc)): for i in range(len(cheb.fc)):
dense_cheb.fc[i].weight.set_data( dense_cheb.fc[i].weight.set_data(cheb.fc[i].weight.data())
cheb.fc[i].weight.data())
if cheb.bias is not None: if cheb.bias is not None:
dense_cheb.bias.set_data( dense_cheb.bias.set_data(cheb.bias.data())
cheb.bias.data())
feat = F.randn((100, 5)) feat = F.randn((100, 5))
out_cheb = cheb(g, feat, [2.0]) out_cheb = cheb(g, feat, [2.0])
out_dense_cheb = dense_cheb(adj, feat, 2.0) out_dense_cheb = dense_cheb(adj, feat, 2.0)
assert F.allclose(out_cheb, out_dense_cheb) assert F.allclose(out_cheb, out_dense_cheb)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('norm_type', ['both', 'right', 'none']) @pytest.mark.parametrize("norm_type", ["both", "right", "none"])
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('out_dim', [1, 2]) "g", get_cases(["homo", "block-bipartite"], exclude=["zero-degree"])
)
@pytest.mark.parametrize("out_dim", [1, 2])
def test_dense_graph_conv(idtype, g, norm_type, out_dim): def test_dense_graph_conv(idtype, g, norm_type, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype('default') adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype("default")
conv = nn.GraphConv(5, out_dim, norm=norm_type, bias=True) conv = nn.GraphConv(5, out_dim, norm=norm_type, bias=True)
dense_conv = nn.DenseGraphConv(5, out_dim, norm=norm_type, bias=True) dense_conv = nn.DenseGraphConv(5, out_dim, norm=norm_type, bias=True)
conv.initialize(ctx=ctx) conv.initialize(ctx=ctx)
dense_conv.initialize(ctx=ctx) dense_conv.initialize(ctx=ctx)
dense_conv.weight.set_data( dense_conv.weight.set_data(conv.weight.data())
conv.weight.data()) dense_conv.bias.set_data(conv.bias.data())
dense_conv.bias.set_data(
conv.bias.data())
feat = F.randn((g.number_of_src_nodes(), 5)) feat = F.randn((g.number_of_src_nodes(), 5))
out_conv = conv(g, feat) out_conv = conv(g, feat)
out_dense_conv = dense_conv(adj, feat) out_dense_conv = dense_conv(adj, feat)
assert F.allclose(out_conv, out_dense_conv) assert F.allclose(out_conv, out_dense_conv)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'bipartite', 'block-bipartite'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('out_dim', [1, 2]) "g", get_cases(["homo", "bipartite", "block-bipartite"])
)
@pytest.mark.parametrize("out_dim", [1, 2])
def test_dense_sage_conv(idtype, g, out_dim): def test_dense_sage_conv(idtype, g, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype('default') adj = g.adjacency_matrix(transpose=True, ctx=ctx).tostype("default")
sage = nn.SAGEConv(5, out_dim, 'gcn') sage = nn.SAGEConv(5, out_dim, "gcn")
dense_sage = nn.DenseSAGEConv(5, out_dim) dense_sage = nn.DenseSAGEConv(5, out_dim)
sage.initialize(ctx=ctx) sage.initialize(ctx=ctx)
dense_sage.initialize(ctx=ctx) dense_sage.initialize(ctx=ctx)
dense_sage.fc.weight.set_data( dense_sage.fc.weight.set_data(sage.fc_neigh.weight.data())
sage.fc_neigh.weight.data()) dense_sage.fc.bias.set_data(sage.fc_neigh.bias.data())
dense_sage.fc.bias.set_data(
sage.fc_neigh.bias.data())
if len(g.ntypes) == 2: if len(g.ntypes) == 2:
feat = ( feat = (
F.randn((g.number_of_src_nodes(), 5)), F.randn((g.number_of_src_nodes(), 5)),
F.randn((g.number_of_dst_nodes(), 5)) F.randn((g.number_of_dst_nodes(), 5)),
) )
else: else:
feat = F.randn((g.number_of_nodes(), 5)) feat = F.randn((g.number_of_nodes(), 5))
...@@ -383,9 +427,12 @@ def test_dense_sage_conv(idtype, g, out_dim): ...@@ -383,9 +427,12 @@ def test_dense_sage_conv(idtype, g, out_dim):
out_dense_sage = dense_sage(adj, feat) out_dense_sage = dense_sage(adj, feat)
assert F.allclose(out_sage, out_dense_sage) assert F.allclose(out_sage, out_dense_sage)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize(
@pytest.mark.parametrize('out_dim', [1, 2]) "g", get_cases(["homo", "block-bipartite"], exclude=["zero-degree"])
)
@pytest.mark.parametrize("out_dim", [1, 2])
def test_edge_conv(g, idtype, out_dim): def test_edge_conv(g, idtype, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -397,9 +444,10 @@ def test_edge_conv(g, idtype, out_dim): ...@@ -397,9 +444,10 @@ def test_edge_conv(g, idtype, out_dim):
h1 = edge_conv(g, h0) h1 = edge_conv(g, h0)
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize("g", get_cases(["bipartite"], exclude=["zero-degree"]))
@pytest.mark.parametrize('out_dim', [1, 2]) @pytest.mark.parametrize("out_dim", [1, 2])
def test_edge_conv_bi(g, idtype, out_dim): def test_edge_conv_bi(g, idtype, out_dim):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -412,9 +460,10 @@ def test_edge_conv_bi(g, idtype, out_dim): ...@@ -412,9 +460,10 @@ def test_edge_conv_bi(g, idtype, out_dim):
h1 = edge_conv(g, (h0, x0)) h1 = edge_conv(g, (h0, x0))
assert h1.shape == (g.number_of_dst_nodes(), out_dim) assert h1.shape == (g.number_of_dst_nodes(), out_dim)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize("g", get_cases(["homo", "block-bipartite"]))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize("aggregator_type", ["mean", "max", "sum"])
def test_gin_conv(g, idtype, aggregator_type): def test_gin_conv(g, idtype, aggregator_type):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -428,9 +477,10 @@ def test_gin_conv(g, idtype, aggregator_type): ...@@ -428,9 +477,10 @@ def test_gin_conv(g, idtype, aggregator_type):
h = gin_conv(g, feat) h = gin_conv(g, feat)
assert h.shape == (g.number_of_dst_nodes(), 5) assert h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize("g", get_cases(["bipartite"]))
@pytest.mark.parametrize('aggregator_type', ['mean', 'max', 'sum']) @pytest.mark.parametrize("aggregator_type", ["mean", "max", "sum"])
def test_gin_conv_bi(g, idtype, aggregator_type): def test_gin_conv_bi(g, idtype, aggregator_type):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -440,29 +490,35 @@ def test_gin_conv_bi(g, idtype, aggregator_type): ...@@ -440,29 +490,35 @@ def test_gin_conv_bi(g, idtype, aggregator_type):
print(gin_conv) print(gin_conv)
# test #2: bipartite # test #2: bipartite
feat = (F.randn((g.number_of_src_nodes(), 5)), F.randn((g.number_of_dst_nodes(), 5))) feat = (
F.randn((g.number_of_src_nodes(), 5)),
F.randn((g.number_of_dst_nodes(), 5)),
)
h = gin_conv(g, feat) h = gin_conv(g, feat)
return h.shape == (g.number_of_dst_nodes(), 5) return h.shape == (g.number_of_dst_nodes(), 5)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize(
"g", get_cases(["homo", "block-bipartite"], exclude=["zero-degree"])
)
def test_gmm_conv(g, idtype): def test_gmm_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
gmm_conv = nn.GMMConv(5, 2, 5, 3, 'max') gmm_conv = nn.GMMConv(5, 2, 5, 3, "max")
gmm_conv.initialize(ctx=ctx) gmm_conv.initialize(ctx=ctx)
h0 = F.randn((g.number_of_src_nodes(), 5)) h0 = F.randn((g.number_of_src_nodes(), 5))
pseudo = F.randn((g.number_of_edges(), 5)) pseudo = F.randn((g.number_of_edges(), 5))
h1 = gmm_conv(g, h0, pseudo) h1 = gmm_conv(g, h0, pseudo)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'], exclude=['zero-degree'])) @pytest.mark.parametrize("g", get_cases(["bipartite"], exclude=["zero-degree"]))
def test_gmm_conv_bi(g, idtype): def test_gmm_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
gmm_conv = nn.GMMConv((5, 4), 2, 5, 3, 'max') gmm_conv = nn.GMMConv((5, 4), 2, 5, 3, "max")
gmm_conv.initialize(ctx=ctx) gmm_conv.initialize(ctx=ctx)
# test #1: basic # test #1: basic
h0 = F.randn((g.number_of_src_nodes(), 5)) h0 = F.randn((g.number_of_src_nodes(), 5))
...@@ -471,12 +527,13 @@ def test_gmm_conv_bi(g, idtype): ...@@ -471,12 +527,13 @@ def test_gmm_conv_bi(g, idtype):
h1 = gmm_conv(g, (h0, hd), pseudo) h1 = gmm_conv(g, (h0, hd), pseudo)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'])) @pytest.mark.parametrize("g", get_cases(["homo", "block-bipartite"]))
def test_nn_conv(g, idtype): def test_nn_conv(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
nn_conv = nn.NNConv(5, 2, gluon.nn.Embedding(3, 5 * 2), 'max') nn_conv = nn.NNConv(5, 2, gluon.nn.Embedding(3, 5 * 2), "max")
nn_conv.initialize(ctx=ctx) nn_conv.initialize(ctx=ctx)
# test #1: basic # test #1: basic
h0 = F.randn((g.number_of_src_nodes(), 5)) h0 = F.randn((g.number_of_src_nodes(), 5))
...@@ -484,12 +541,13 @@ def test_nn_conv(g, idtype): ...@@ -484,12 +541,13 @@ def test_nn_conv(g, idtype):
h1 = nn_conv(g, h0, etypes) h1 = nn_conv(g, h0, etypes)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['bipartite'])) @pytest.mark.parametrize("g", get_cases(["bipartite"]))
def test_nn_conv_bi(g, idtype): def test_nn_conv_bi(g, idtype):
g = g.astype(idtype).to(F.ctx()) g = g.astype(idtype).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
nn_conv = nn.NNConv((5, 4), 2, gluon.nn.Embedding(3, 5 * 2), 'max') nn_conv = nn.NNConv((5, 4), 2, gluon.nn.Embedding(3, 5 * 2), "max")
nn_conv.initialize(ctx=ctx) nn_conv.initialize(ctx=ctx)
# test #1: basic # test #1: basic
h0 = F.randn((g.number_of_src_nodes(), 5)) h0 = F.randn((g.number_of_src_nodes(), 5))
...@@ -498,7 +556,8 @@ def test_nn_conv_bi(g, idtype): ...@@ -498,7 +556,8 @@ def test_nn_conv_bi(g, idtype):
h1 = nn_conv(g, (h0, hd), etypes) h1 = nn_conv(g, (h0, hd), etypes)
assert h1.shape == (g.number_of_dst_nodes(), 2) assert h1.shape == (g.number_of_dst_nodes(), 2)
@pytest.mark.parametrize('out_dim', [1, 2])
@pytest.mark.parametrize("out_dim", [1, 2])
def test_sg_conv(out_dim): def test_sg_conv(out_dim):
g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx()) g = dgl.from_networkx(nx.erdos_renyi_graph(20, 0.3)).to(F.ctx())
g = dgl.add_self_loop(g) g = dgl.add_self_loop(g)
...@@ -513,11 +572,12 @@ def test_sg_conv(out_dim): ...@@ -513,11 +572,12 @@ def test_sg_conv(out_dim):
h1 = sgc(g, h0) h1 = sgc(g, h0)
assert h1.shape == (g.number_of_nodes(), out_dim) assert h1.shape == (g.number_of_nodes(), out_dim)
def test_set2set(): def test_set2set():
g = dgl.from_networkx(nx.path_graph(10)).to(F.ctx()) g = dgl.from_networkx(nx.path_graph(10)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
s2s = nn.Set2Set(5, 3, 3) # hidden size 5, 3 iters, 3 layers s2s = nn.Set2Set(5, 3, 3) # hidden size 5, 3 iters, 3 layers
s2s.initialize(ctx=ctx) s2s.initialize(ctx=ctx)
print(s2s) print(s2s)
...@@ -532,6 +592,7 @@ def test_set2set(): ...@@ -532,6 +592,7 @@ def test_set2set():
h1 = s2s(bg, h0) h1 = s2s(bg, h0)
assert h1.shape[0] == 3 and h1.shape[1] == 10 and h1.ndim == 2 assert h1.shape[0] == 3 and h1.shape[1] == 10 and h1.ndim == 2
def test_glob_att_pool(): def test_glob_att_pool():
g = dgl.from_networkx(nx.path_graph(10)).to(F.ctx()) g = dgl.from_networkx(nx.path_graph(10)).to(F.ctx())
ctx = F.ctx() ctx = F.ctx()
...@@ -550,13 +611,14 @@ def test_glob_att_pool(): ...@@ -550,13 +611,14 @@ def test_glob_att_pool():
h1 = gap(bg, h0) h1 = gap(bg, h0)
assert h1.shape[0] == 4 and h1.shape[1] == 10 and h1.ndim == 2 assert h1.shape[0] == 4 and h1.shape[1] == 10 and h1.ndim == 2
def test_simple_pool(): def test_simple_pool():
g = dgl.from_networkx(nx.path_graph(15)).to(F.ctx()) g = dgl.from_networkx(nx.path_graph(15)).to(F.ctx())
sum_pool = nn.SumPooling() sum_pool = nn.SumPooling()
avg_pool = nn.AvgPooling() avg_pool = nn.AvgPooling()
max_pool = nn.MaxPooling() max_pool = nn.MaxPooling()
sort_pool = nn.SortPooling(10) # k = 10 sort_pool = nn.SortPooling(10) # k = 10
print(sum_pool, avg_pool, max_pool, sort_pool) print(sum_pool, avg_pool, max_pool, sort_pool)
# test#1: basic # test#1: basic
...@@ -575,33 +637,43 @@ def test_simple_pool(): ...@@ -575,33 +637,43 @@ def test_simple_pool():
bg = dgl.batch([g, g_, g, g_, g]) bg = dgl.batch([g, g_, g, g_, g])
h0 = F.randn((bg.number_of_nodes(), 5)) h0 = F.randn((bg.number_of_nodes(), 5))
h1 = sum_pool(bg, h0) h1 = sum_pool(bg, h0)
truth = mx.nd.stack(F.sum(h0[:15], 0), truth = mx.nd.stack(
F.sum(h0[15:20], 0), F.sum(h0[:15], 0),
F.sum(h0[20:35], 0), F.sum(h0[15:20], 0),
F.sum(h0[35:40], 0), F.sum(h0[20:35], 0),
F.sum(h0[40:55], 0), axis=0) F.sum(h0[35:40], 0),
F.sum(h0[40:55], 0),
axis=0,
)
check_close(h1, truth) check_close(h1, truth)
h1 = avg_pool(bg, h0) h1 = avg_pool(bg, h0)
truth = mx.nd.stack(F.mean(h0[:15], 0), truth = mx.nd.stack(
F.mean(h0[15:20], 0), F.mean(h0[:15], 0),
F.mean(h0[20:35], 0), F.mean(h0[15:20], 0),
F.mean(h0[35:40], 0), F.mean(h0[20:35], 0),
F.mean(h0[40:55], 0), axis=0) F.mean(h0[35:40], 0),
F.mean(h0[40:55], 0),
axis=0,
)
check_close(h1, truth) check_close(h1, truth)
h1 = max_pool(bg, h0) h1 = max_pool(bg, h0)
truth = mx.nd.stack(F.max(h0[:15], 0), truth = mx.nd.stack(
F.max(h0[15:20], 0), F.max(h0[:15], 0),
F.max(h0[20:35], 0), F.max(h0[15:20], 0),
F.max(h0[35:40], 0), F.max(h0[20:35], 0),
F.max(h0[40:55], 0), axis=0) F.max(h0[35:40], 0),
F.max(h0[40:55], 0),
axis=0,
)
check_close(h1, truth) check_close(h1, truth)
h1 = sort_pool(bg, h0) h1 = sort_pool(bg, h0)
assert h1.shape[0] == 5 and h1.shape[1] == 10 * 5 and h1.ndim == 2 assert h1.shape[0] == 5 and h1.shape[1] == 10 * 5 and h1.ndim == 2
@pytest.mark.parametrize('O', [1, 2, 8])
@pytest.mark.parametrize("O", [1, 2, 8])
def test_rgcn(O): def test_rgcn(O):
ctx = F.ctx() ctx = F.ctx()
etype = [] etype = []
...@@ -654,6 +726,7 @@ def test_rgcn(O): ...@@ -654,6 +726,7 @@ def test_rgcn(O):
h_new = rgc_basis(g, h, r) h_new = rgc_basis(g, h, r)
assert list(h_new.shape) == [100, O] assert list(h_new.shape) == [100, O]
def test_sequential(): def test_sequential():
ctx = F.ctx() ctx = F.ctx()
# test single graph # test single graph
...@@ -663,11 +736,11 @@ def test_sequential(): ...@@ -663,11 +736,11 @@ def test_sequential():
def forward(self, graph, n_feat, e_feat): def forward(self, graph, n_feat, e_feat):
graph = graph.local_var() graph = graph.local_var()
graph.ndata['h'] = n_feat graph.ndata["h"] = n_feat
graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h')) graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
n_feat += graph.ndata['h'] n_feat += graph.ndata["h"]
graph.apply_edges(fn.u_add_v('h', 'h', 'e')) graph.apply_edges(fn.u_add_v("h", "h", "e"))
e_feat += graph.edata['e'] e_feat += graph.edata["e"]
return n_feat, e_feat return n_feat, e_feat
g = dgl.graph(([], [])).to(F.ctx()) g = dgl.graph(([], [])).to(F.ctx())
...@@ -691,9 +764,9 @@ def test_sequential(): ...@@ -691,9 +764,9 @@ def test_sequential():
def forward(self, graph, n_feat): def forward(self, graph, n_feat):
graph = graph.local_var() graph = graph.local_var()
graph.ndata['h'] = n_feat graph.ndata["h"] = n_feat
graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h')) graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
n_feat += graph.ndata['h'] n_feat += graph.ndata["h"]
return n_feat.reshape(graph.number_of_nodes() // 2, 2, -1).sum(1) return n_feat.reshape(graph.number_of_nodes() // 2, 2, -1).sum(1)
g1 = dgl.from_networkx(nx.erdos_renyi_graph(32, 0.05)).to(F.ctx()) g1 = dgl.from_networkx(nx.erdos_renyi_graph(32, 0.05)).to(F.ctx())
...@@ -709,58 +782,75 @@ def test_sequential(): ...@@ -709,58 +782,75 @@ def test_sequential():
n_feat = net([g1, g2, g3], n_feat) n_feat = net([g1, g2, g3], n_feat)
assert n_feat.shape == (4, 4) assert n_feat.shape == (4, 4)
def myagg(alist, dsttype): def myagg(alist, dsttype):
rst = alist[0] rst = alist[0]
for i in range(1, len(alist)): for i in range(1, len(alist)):
rst = rst + (i + 1) * alist[i] rst = rst + (i + 1) * alist[i]
return rst return rst
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('agg', ['sum', 'max', 'min', 'mean', 'stack', myagg]) @pytest.mark.parametrize("agg", ["sum", "max", "min", "mean", "stack", myagg])
def test_hetero_conv(agg, idtype): def test_hetero_conv(agg, idtype):
g = dgl.heterograph({ g = dgl.heterograph(
('user', 'follows', 'user'): ([0, 0, 2, 1], [1, 2, 1, 3]), {
('user', 'plays', 'game'): ([0, 0, 0, 1, 2], [0, 2, 3, 0, 2]), ("user", "follows", "user"): ([0, 0, 2, 1], [1, 2, 1, 3]),
('store', 'sells', 'game'): ([0, 0, 1, 1], [0, 3, 1, 2])}, ("user", "plays", "game"): ([0, 0, 0, 1, 2], [0, 2, 3, 0, 2]),
idtype=idtype, device=F.ctx()) ("store", "sells", "game"): ([0, 0, 1, 1], [0, 3, 1, 2]),
conv = nn.HeteroGraphConv({ },
'follows': nn.GraphConv(2, 3, allow_zero_in_degree=True), idtype=idtype,
'plays': nn.GraphConv(2, 4, allow_zero_in_degree=True), device=F.ctx(),
'sells': nn.GraphConv(3, 4, allow_zero_in_degree=True)}, )
agg) conv = nn.HeteroGraphConv(
{
"follows": nn.GraphConv(2, 3, allow_zero_in_degree=True),
"plays": nn.GraphConv(2, 4, allow_zero_in_degree=True),
"sells": nn.GraphConv(3, 4, allow_zero_in_degree=True),
},
agg,
)
conv.initialize(ctx=F.ctx()) conv.initialize(ctx=F.ctx())
print(conv) print(conv)
uf = F.randn((4, 2)) uf = F.randn((4, 2))
gf = F.randn((4, 4)) gf = F.randn((4, 4))
sf = F.randn((2, 3)) sf = F.randn((2, 3))
h = conv(g, {'user': uf, 'store': sf, 'game': gf}) h = conv(g, {"user": uf, "store": sf, "game": gf})
assert set(h.keys()) == {'user', 'game'} assert set(h.keys()) == {"user", "game"}
if agg != 'stack': if agg != "stack":
assert h['user'].shape == (4, 3) assert h["user"].shape == (4, 3)
assert h['game'].shape == (4, 4) assert h["game"].shape == (4, 4)
else: else:
assert h['user'].shape == (4, 1, 3) assert h["user"].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4) assert h["game"].shape == (4, 2, 4)
block = dgl.to_block(g.to(F.cpu()), {'user': [0, 1, 2, 3], 'game': [0, 1, 2, 3], 'store': []}).to(F.ctx()) block = dgl.to_block(
h = conv(block, ({'user': uf, 'game': gf, 'store': sf}, {'user': uf, 'game': gf, 'store': sf[0:0]})) g.to(F.cpu()), {"user": [0, 1, 2, 3], "game": [0, 1, 2, 3], "store": []}
assert set(h.keys()) == {'user', 'game'} ).to(F.ctx())
if agg != 'stack': h = conv(
assert h['user'].shape == (4, 3) block,
assert h['game'].shape == (4, 4) (
{"user": uf, "game": gf, "store": sf},
{"user": uf, "game": gf, "store": sf[0:0]},
),
)
assert set(h.keys()) == {"user", "game"}
if agg != "stack":
assert h["user"].shape == (4, 3)
assert h["game"].shape == (4, 4)
else: else:
assert h['user'].shape == (4, 1, 3) assert h["user"].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4) assert h["game"].shape == (4, 2, 4)
h = conv(block, {'user': uf, 'game': gf, 'store': sf}) h = conv(block, {"user": uf, "game": gf, "store": sf})
assert set(h.keys()) == {'user', 'game'} assert set(h.keys()) == {"user", "game"}
if agg != 'stack': if agg != "stack":
assert h['user'].shape == (4, 3) assert h["user"].shape == (4, 3)
assert h['game'].shape == (4, 4) assert h["game"].shape == (4, 4)
else: else:
assert h['user'].shape == (4, 1, 3) assert h["user"].shape == (4, 1, 3)
assert h['game'].shape == (4, 2, 4) assert h["game"].shape == (4, 2, 4)
# test with mod args # test with mod args
class MyMod(mx.gluon.nn.Block): class MyMod(mx.gluon.nn.Block):
...@@ -769,39 +859,46 @@ def test_hetero_conv(agg, idtype): ...@@ -769,39 +859,46 @@ def test_hetero_conv(agg, idtype):
self.carg1 = 0 self.carg1 = 0
self.s1 = s1 self.s1 = s1
self.s2 = s2 self.s2 = s2
def forward(self, g, h, arg1=None): # mxnet does not support kwargs def forward(self, g, h, arg1=None): # mxnet does not support kwargs
if arg1 is not None: if arg1 is not None:
self.carg1 += 1 self.carg1 += 1
return F.zeros((g.number_of_dst_nodes(), self.s2)) return F.zeros((g.number_of_dst_nodes(), self.s2))
mod1 = MyMod(2, 3) mod1 = MyMod(2, 3)
mod2 = MyMod(2, 4) mod2 = MyMod(2, 4)
mod3 = MyMod(3, 4) mod3 = MyMod(3, 4)
conv = nn.HeteroGraphConv({ conv = nn.HeteroGraphConv(
'follows': mod1, {"follows": mod1, "plays": mod2, "sells": mod3}, agg
'plays': mod2, )
'sells': mod3},
agg)
conv.initialize(ctx=F.ctx()) conv.initialize(ctx=F.ctx())
mod_args = {'follows' : (1,), 'plays' : (1,)} mod_args = {"follows": (1,), "plays": (1,)}
h = conv(g, {'user' : uf, 'store' : sf, 'game': gf}, mod_args) h = conv(g, {"user": uf, "store": sf, "game": gf}, mod_args)
assert mod1.carg1 == 1 assert mod1.carg1 == 1
assert mod2.carg1 == 1 assert mod2.carg1 == 1
assert mod3.carg1 == 0 assert mod3.carg1 == 0
#conv on graph without any edges # conv on graph without any edges
for etype in g.etypes: for etype in g.etypes:
g = dgl.remove_edges(g, g.edges(form='eid', etype=etype), etype=etype) g = dgl.remove_edges(g, g.edges(form="eid", etype=etype), etype=etype)
assert g.num_edges() == 0 assert g.num_edges() == 0
h = conv(g, {'user': uf, 'game': gf, 'store': sf}) h = conv(g, {"user": uf, "game": gf, "store": sf})
assert set(h.keys()) == {'user', 'game'} assert set(h.keys()) == {"user", "game"}
block = dgl.to_block(g.to(F.cpu()), {'user': [0, 1, 2, 3], 'game': [ block = dgl.to_block(
0, 1, 2, 3], 'store': []}).to(F.ctx()) g.to(F.cpu()), {"user": [0, 1, 2, 3], "game": [0, 1, 2, 3], "store": []}
h = conv(block, ({'user': uf, 'game': gf, 'store': sf}, ).to(F.ctx())
{'user': uf, 'game': gf, 'store': sf[0:0]})) h = conv(
assert set(h.keys()) == {'user', 'game'} block,
(
if __name__ == '__main__': {"user": uf, "game": gf, "store": sf},
{"user": uf, "game": gf, "store": sf[0:0]},
),
)
assert set(h.keys()) == {"user", "game"}
if __name__ == "__main__":
test_graph_conv() test_graph_conv()
test_gat_conv() test_gat_conv()
test_sage_conv() test_sage_conv()
......
tests/pytorch/mock_sparse/test_elementwise_op_sp.py
View file @ 89a4cc4d
import operator
import numpy as np import numpy as np
import pytest import pytest
import dgl
import torch import torch
import operator
import dgl
from dgl.mock_sparse import SparseMatrix, diag from dgl.mock_sparse import SparseMatrix, diag
parametrize_idtype = pytest.mark.parametrize( parametrize_idtype = pytest.mark.parametrize(
......
tests/pytorch/test_dist_optim.py
View file @ 89a4cc4d
import os import os
os.environ['OMP_NUM_THREADS'] = '1'
import dgl os.environ["OMP_NUM_THREADS"] = "1"
import multiprocessing as mp
import pickle
import random
import socket
import sys import sys
import numpy as np
import time import time
import socket import unittest
from scipy import sparse as spsp
import backend as F
import numpy as np
import torch as th import torch as th
from scipy import sparse as spsp
from dgl.distributed import DistGraphServer, DistGraph import dgl
from dgl.distributed import partition_graph, load_partition_book
import multiprocessing as mp
from dgl import function as fn from dgl import function as fn
import backend as F from dgl.distributed import (
import unittest DistEmbedding,
import pickle DistGraph,
import random DistGraphServer,
from dgl.distributed import DistEmbedding load_partition_book,
partition_graph,
)
from dgl.distributed.optim import SparseAdagrad, SparseAdam from dgl.distributed.optim import SparseAdagrad, SparseAdam
def create_random_graph(n): def create_random_graph(n):
arr = (spsp.random(n, n, density=0.001, format='coo', random_state=100) != 0).astype(np.int64) arr = (
spsp.random(n, n, density=0.001, format="coo", random_state=100) != 0
).astype(np.int64)
return dgl.from_scipy(arr) return dgl.from_scipy(arr)
def get_local_usable_addr(): def get_local_usable_addr():
"""Get local usable IP and port """Get local usable IP and port
...@@ -34,10 +44,10 @@ def get_local_usable_addr(): ...@@ -34,10 +44,10 @@ def get_local_usable_addr():
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try: try:
# doesn't even have to be reachable # doesn't even have to be reachable
sock.connect(('10.255.255.255', 1)) sock.connect(("10.255.255.255", 1))
ip_addr = sock.getsockname()[0] ip_addr = sock.getsockname()[0]
except ValueError: except ValueError:
ip_addr = '127.0.0.1' ip_addr = "127.0.0.1"
finally: finally:
sock.close() sock.close()
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
...@@ -46,40 +56,62 @@ def get_local_usable_addr(): ...@@ -46,40 +56,62 @@ def get_local_usable_addr():
port = sock.getsockname()[1] port = sock.getsockname()[1]
sock.close() sock.close()
return ip_addr + ' ' + str(port) return ip_addr + " " + str(port)
def prepare_dist(): def prepare_dist():
ip_config = open("optim_ip_config.txt", "w") ip_config = open("optim_ip_config.txt", "w")
ip_addr = get_local_usable_addr() ip_addr = get_local_usable_addr()
ip_config.write('{}\n'.format(ip_addr)) ip_config.write("{}\n".format(ip_addr))
ip_config.close() ip_config.close()
def run_server(graph_name, server_id, server_count, num_clients, shared_mem): def run_server(graph_name, server_id, server_count, num_clients, shared_mem):
g = DistGraphServer(server_id, "optim_ip_config.txt", num_clients, server_count, g = DistGraphServer(
'/tmp/dist_graph/{}.json'.format(graph_name), server_id,
disable_shared_mem=not shared_mem) "optim_ip_config.txt",
print('start server', server_id) num_clients,
server_count,
"/tmp/dist_graph/{}.json".format(graph_name),
disable_shared_mem=not shared_mem,
)
print("start server", server_id)
g.start() g.start()
def initializer(shape, dtype): def initializer(shape, dtype):
arr = th.zeros(shape, dtype=dtype) arr = th.zeros(shape, dtype=dtype)
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(arr, 0, 1.0) th.nn.init.uniform_(arr, 0, 1.0)
return arr return arr
def run_client(graph_name, cli_id, part_id, server_count): def run_client(graph_name, cli_id, part_id, server_count):
device=F.ctx() device = F.ctx()
time.sleep(5) time.sleep(5)
os.environ['DGL_NUM_SERVER'] = str(server_count) os.environ["DGL_NUM_SERVER"] = str(server_count)
dgl.distributed.initialize("optim_ip_config.txt") dgl.distributed.initialize("optim_ip_config.txt")
gpb, graph_name, _, _ = load_partition_book('/tmp/dist_graph/{}.json'.format(graph_name), gpb, graph_name, _, _ = load_partition_book(
part_id, None) "/tmp/dist_graph/{}.json".format(graph_name), part_id, None
)
g = DistGraph(graph_name, gpb=gpb) g = DistGraph(graph_name, gpb=gpb)
policy = dgl.distributed.PartitionPolicy('node', g.get_partition_book()) policy = dgl.distributed.PartitionPolicy("node", g.get_partition_book())
num_nodes = g.number_of_nodes() num_nodes = g.number_of_nodes()
emb_dim = 4 emb_dim = 4
dgl_emb = DistEmbedding(num_nodes, emb_dim, name='optim', init_func=initializer, part_policy=policy) dgl_emb = DistEmbedding(
dgl_emb_zero = DistEmbedding(num_nodes, emb_dim, name='optim-zero', init_func=initializer, part_policy=policy) num_nodes,
emb_dim,
name="optim",
init_func=initializer,
part_policy=policy,
)
dgl_emb_zero = DistEmbedding(
num_nodes,
emb_dim,
name="optim-zero",
init_func=initializer,
part_policy=policy,
)
dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01) dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
dgl_adam._world_size = 1 dgl_adam._world_size = 1
dgl_adam._rank = 0 dgl_adam._rank = 0
...@@ -91,11 +123,13 @@ def run_client(graph_name, cli_id, part_id, server_count): ...@@ -91,11 +123,13 @@ def run_client(graph_name, cli_id, part_id, server_count):
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(torch_emb_zero.weight, 0, 1.0) th.nn.init.uniform_(torch_emb_zero.weight, 0, 1.0)
torch_adam = th.optim.SparseAdam( torch_adam = th.optim.SparseAdam(
list(torch_emb.parameters()) + list(torch_emb_zero.parameters()), lr=0.01) list(torch_emb.parameters()) + list(torch_emb_zero.parameters()),
lr=0.01,
)
labels = th.ones((4,)).long() labels = th.ones((4,)).long()
idx = th.randint(0, num_nodes, size=(4,)) idx = th.randint(0, num_nodes, size=(4,))
dgl_value = dgl_emb(idx, device).to(th.device('cpu')) dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
torch_value = torch_emb(idx) torch_value = torch_emb(idx)
torch_adam.zero_grad() torch_adam.zero_grad()
torch_loss = th.nn.functional.cross_entropy(torch_value, labels) torch_loss = th.nn.functional.cross_entropy(torch_value, labels)
...@@ -107,7 +141,10 @@ def run_client(graph_name, cli_id, part_id, server_count): ...@@ -107,7 +141,10 @@ def run_client(graph_name, cli_id, part_id, server_count):
dgl_loss.backward() dgl_loss.backward()
dgl_adam.step() dgl_adam.step()
assert F.allclose(dgl_emb.weight[0 : num_nodes//2], torch_emb.weight[0 : num_nodes//2]) assert F.allclose(
dgl_emb.weight[0 : num_nodes // 2], torch_emb.weight[0 : num_nodes // 2]
)
def check_sparse_adam(num_trainer=1, shared_mem=True): def check_sparse_adam(num_trainer=1, shared_mem=True):
prepare_dist() prepare_dist()
...@@ -116,23 +153,27 @@ def check_sparse_adam(num_trainer=1, shared_mem=True): ...@@ -116,23 +153,27 @@ def check_sparse_adam(num_trainer=1, shared_mem=True):
num_clients = num_trainer num_clients = num_trainer
num_parts = 1 num_parts = 1
graph_name = 'dist_graph_test' graph_name = "dist_graph_test"
partition_graph(g, graph_name, num_parts, '/tmp/dist_graph') partition_graph(g, graph_name, num_parts, "/tmp/dist_graph")
# let's just test on one partition for now. # let's just test on one partition for now.
# We cannot run multiple servers and clients on the same machine. # We cannot run multiple servers and clients on the same machine.
serv_ps = [] serv_ps = []
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for serv_id in range(num_servers): for serv_id in range(num_servers):
p = ctx.Process(target=run_server, args=(graph_name, serv_id, num_servers, p = ctx.Process(
num_clients, shared_mem)) target=run_server,
args=(graph_name, serv_id, num_servers, num_clients, shared_mem),
)
serv_ps.append(p) serv_ps.append(p)
p.start() p.start()
cli_ps = [] cli_ps = []
for cli_id in range(num_clients): for cli_id in range(num_clients):
print('start client', cli_id) print("start client", cli_id)
p = ctx.Process(target=run_client, args=(graph_name, cli_id, 0, num_servers)) p = ctx.Process(
target=run_client, args=(graph_name, cli_id, 0, num_servers)
)
p.start() p.start()
cli_ps.append(p) cli_ps.append(p)
...@@ -142,12 +183,14 @@ def check_sparse_adam(num_trainer=1, shared_mem=True): ...@@ -142,12 +183,14 @@ def check_sparse_adam(num_trainer=1, shared_mem=True):
for p in serv_ps: for p in serv_ps:
p.join() p.join()
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
def test_sparse_opt(): def test_sparse_opt():
os.environ['DGL_DIST_MODE'] = 'distributed' os.environ["DGL_DIST_MODE"] = "distributed"
check_sparse_adam(1, True) check_sparse_adam(1, True)
check_sparse_adam(1, False) check_sparse_adam(1, False)
if __name__ == '__main__':
os.makedirs('/tmp/dist_graph', exist_ok=True) if __name__ == "__main__":
test_sparse_opt() os.makedirs("/tmp/dist_graph", exist_ok=True)
\ No newline at end of file test_sparse_opt()
tests/pytorch/test_geometry.py
View file @ 89a4cc4d
import backend as F import backend as F
import dgl.nn
import dgl
import numpy as np import numpy as np
import pytest import pytest
import torch as th import torch as th
from dgl import DGLError
from dgl.base import DGLWarning
from dgl.geometry import neighbor_matching, farthest_point_sampler
from test_utils import parametrize_idtype from test_utils import parametrize_idtype
from test_utils.graph_cases import get_cases from test_utils.graph_cases import get_cases
import dgl
import dgl.nn
from dgl import DGLError
from dgl.base import DGLWarning
from dgl.geometry import farthest_point_sampler, neighbor_matching
def test_fps(): def test_fps():
N = 1000 N = 1000
batch_size = 5 batch_size = 5
sample_points = 10 sample_points = 10
x = th.tensor(np.random.uniform(size=(batch_size, int(N/batch_size), 3))) x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
ctx = F.ctx() ctx = F.ctx()
if F.gpu_ctx(): if F.gpu_ctx():
x = x.to(ctx) x = x.to(ctx)
...@@ -29,17 +30,18 @@ def test_fps_start_idx(): ...@@ -29,17 +30,18 @@ def test_fps_start_idx():
N = 1000 N = 1000
batch_size = 5 batch_size = 5
sample_points = 10 sample_points = 10
x = th.tensor(np.random.uniform(size=(batch_size, int(N/batch_size), 3))) x = th.tensor(np.random.uniform(size=(batch_size, int(N / batch_size), 3)))
ctx = F.ctx() ctx = F.ctx()
if F.gpu_ctx(): if F.gpu_ctx():
x = x.to(ctx) x = x.to(ctx)
res = farthest_point_sampler(x, sample_points, start_idx=0) res = farthest_point_sampler(x, sample_points, start_idx=0)
assert th.any(res[:, 0] == 0) assert th.any(res[:, 0] == 0)
def _test_knn_common(device, algorithm, dist, exclude_self): def _test_knn_common(device, algorithm, dist, exclude_self):
x = th.randn(8, 3).to(device) x = th.randn(8, 3).to(device)
kg = dgl.nn.KNNGraph(3) kg = dgl.nn.KNNGraph(3)
if dist == 'euclidean': if dist == "euclidean":
d = th.cdist(x, x).to(F.cpu()) d = th.cdist(x, x).to(F.cpu())
else: else:
x = x + th.randn(1).item() x = x + th.randn(1).item()
...@@ -55,7 +57,14 @@ def _test_knn_common(device, algorithm, dist, exclude_self): ...@@ -55,7 +57,14 @@ def _test_knn_common(device, algorithm, dist, exclude_self):
assert len(src) == k assert len(src) == k
if check_indices: if check_indices:
i = v - start i = v - start
src_ans = set(th.topk(d[start:end, start:end][i], k + (1 if exclude_self else 0), largest=False)[1].numpy() + start) src_ans = set(
th.topk(
d[start:end, start:end][i],
k + (1 if exclude_self else 0),
largest=False,
)[1].numpy()
+ start
)
if exclude_self: if exclude_self:
# remove self # remove self
src_ans.remove(v) src_ans.remove(v)
...@@ -63,7 +72,9 @@ def _test_knn_common(device, algorithm, dist, exclude_self): ...@@ -63,7 +72,9 @@ def _test_knn_common(device, algorithm, dist, exclude_self):
def check_batch(g, k, expected_batch_info): def check_batch(g, k, expected_batch_info):
assert F.array_equal(g.batch_num_nodes(), F.tensor(expected_batch_info)) assert F.array_equal(g.batch_num_nodes(), F.tensor(expected_batch_info))
assert F.array_equal(g.batch_num_edges(), k*F.tensor(expected_batch_info)) assert F.array_equal(
g.batch_num_edges(), k * F.tensor(expected_batch_info)
)
# check knn with 2d input # check knn with 2d input
g = kg(x, algorithm, dist, exclude_self) g = kg(x, algorithm, dist, exclude_self)
...@@ -145,23 +156,27 @@ def _test_knn_common(device, algorithm, dist, exclude_self): ...@@ -145,23 +156,27 @@ def _test_knn_common(device, algorithm, dist, exclude_self):
kg = dgl.nn.SegmentedKNNGraph(3) kg = dgl.nn.SegmentedKNNGraph(3)
g = kg(x, [4, 7, 5, 4], algorithm, dist, exclude_self) g = kg(x, [4, 7, 5, 4], algorithm, dist, exclude_self)
# different algorithms may break the tie differently, so don't check the indices # different algorithms may break the tie differently, so don't check the indices
check_knn(g, x, 0, 4, 3, exclude_self, False) check_knn(g, x, 0, 4, 3, exclude_self, False)
check_knn(g, x, 4, 11, 3, exclude_self, False) check_knn(g, x, 4, 11, 3, exclude_self, False)
check_knn(g, x, 11, 16, 3, exclude_self, False) check_knn(g, x, 11, 16, 3, exclude_self, False)
check_knn(g, x, 16, 20, 3, exclude_self, False) check_knn(g, x, 16, 20, 3, exclude_self, False)
check_batch(g, 3, [4, 7, 5, 4]) check_batch(g, 3, [4, 7, 5, 4])
@pytest.mark.parametrize('algorithm', ['bruteforce-blas', 'bruteforce', 'kd-tree']) @pytest.mark.parametrize(
@pytest.mark.parametrize('dist', ['euclidean', 'cosine']) "algorithm", ["bruteforce-blas", "bruteforce", "kd-tree"]
@pytest.mark.parametrize('exclude_self', [False, True]) )
@pytest.mark.parametrize("dist", ["euclidean", "cosine"])
@pytest.mark.parametrize("exclude_self", [False, True])
def test_knn_cpu(algorithm, dist, exclude_self): def test_knn_cpu(algorithm, dist, exclude_self):
_test_knn_common(F.cpu(), algorithm, dist, exclude_self) _test_knn_common(F.cpu(), algorithm, dist, exclude_self)
@pytest.mark.parametrize('algorithm', ['bruteforce-blas', 'bruteforce', 'bruteforce-sharemem']) @pytest.mark.parametrize(
@pytest.mark.parametrize('dist', ['euclidean', 'cosine']) "algorithm", ["bruteforce-blas", "bruteforce", "bruteforce-sharemem"]
@pytest.mark.parametrize('exclude_self', [False, True]) )
@pytest.mark.parametrize("dist", ["euclidean", "cosine"])
@pytest.mark.parametrize("exclude_self", [False, True])
def test_knn_cuda(algorithm, dist, exclude_self): def test_knn_cuda(algorithm, dist, exclude_self):
if not th.cuda.is_available(): if not th.cuda.is_available():
return return
...@@ -169,9 +184,9 @@ def test_knn_cuda(algorithm, dist, exclude_self): ...@@ -169,9 +184,9 @@ def test_knn_cuda(algorithm, dist, exclude_self):
@parametrize_idtype @parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['dglgraph'])) @pytest.mark.parametrize("g", get_cases(["homo"], exclude=["dglgraph"]))
@pytest.mark.parametrize('weight', [True, False]) @pytest.mark.parametrize("weight", [True, False])
@pytest.mark.parametrize('relabel', [True, False]) @pytest.mark.parametrize("relabel", [True, False])
def test_edge_coarsening(idtype, g, weight, relabel): def test_edge_coarsening(idtype, g, weight, relabel):
num_nodes = g.num_nodes() num_nodes = g.num_nodes()
g = dgl.to_bidirected(g) g = dgl.to_bidirected(g)
...@@ -205,7 +220,7 @@ def test_edge_coarsening(idtype, g, weight, relabel): ...@@ -205,7 +220,7 @@ def test_edge_coarsening(idtype, g, weight, relabel):
assert g.has_edges_between(u, v) assert g.has_edges_between(u, v)
if __name__ == '__main__': if __name__ == "__main__":
test_fps() test_fps()
test_fps_start_idx() test_fps_start_idx()
test_knn() test_knn()
tests/pytorch/test_ipc.py
View file @ 89a4cc4d
import dgl
import torch as th
import torch.multiprocessing as mp
import os import os
import unittest import unittest
import torch as th
import torch.multiprocessing as mp
import dgl
def sub_ipc(g): def sub_ipc(g):
print(g) print(g)
return g return g
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
def test_torch_ipc(): def test_torch_ipc():
g = dgl.graph(([0, 1, 2], [1, 2, 3])) g = dgl.graph(([0, 1, 2], [1, 2, 3]))
ctx = mp.get_context("spawn") ctx = mp.get_context("spawn")
p = ctx.Process(target=sub_ipc, args=(g, )) p = ctx.Process(target=sub_ipc, args=(g,))
p.start() p.start()
p.join() p.join()
if __name__ == "__main__": if __name__ == "__main__":
test_torch_ipc() test_torch_ipc()
\ No newline at end of file
tests/pytorch/test_optim.py
View file @ 89a4cc4d
import os
import time import time
import torch.multiprocessing as mp import unittest
import unittest, os
import pytest
import torch as th
import backend as F import backend as F
import pytest
import torch as th
import torch.multiprocessing as mp
from dgl.nn import NodeEmbedding from dgl.nn import NodeEmbedding
from dgl.optim import SparseAdam, SparseAdagrad from dgl.optim import SparseAdagrad, SparseAdam
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@pytest.mark.parametrize('emb_dim', [1, 4, 101, 1024]) @pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
def test_sparse_adam(emb_dim): def test_sparse_adam(emb_dim):
num_embs = 10 num_embs = 10
device=F.ctx() device = F.ctx()
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test') dgl_emb = NodeEmbedding(num_embs, emb_dim, "test")
torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(torch_emb.weight, 0, 1.0) th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
...@@ -26,7 +28,7 @@ def test_sparse_adam(emb_dim): ...@@ -26,7 +28,7 @@ def test_sparse_adam(emb_dim):
# first step # first step
idx = th.randint(0, num_embs, size=(4,)) idx = th.randint(0, num_embs, size=(4,))
dgl_value = dgl_emb(idx, device).to(th.device('cpu')) dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
torch_value = torch_emb(idx) torch_value = torch_emb(idx)
labels = th.zeros((4,)).long() labels = th.zeros((4,)).long()
print("dgl_value = {}".format(dgl_value)) print("dgl_value = {}".format(dgl_value))
...@@ -47,17 +49,18 @@ def test_sparse_adam(emb_dim): ...@@ -47,17 +49,18 @@ def test_sparse_adam(emb_dim):
# Pytorch sparseAdam maintains a global step # Pytorch sparseAdam maintains a global step
# DGL sparseAdam use a per embedding step # DGL sparseAdam use a per embedding step
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@pytest.mark.parametrize('use_uva', [False, True, None]) @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@pytest.mark.parametrize('emb_dim', [1, 4, 101, 1024]) @pytest.mark.parametrize("use_uva", [False, True, None])
@pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
def test_sparse_adam_uva(use_uva, emb_dim): def test_sparse_adam_uva(use_uva, emb_dim):
if F.ctx().type == 'cpu' and use_uva == True: if F.ctx().type == "cpu" and use_uva == True:
# we want to only test values of False and None when not using GPU # we want to only test values of False and None when not using GPU
pytest.skip("UVA cannot be used without GPUs.") pytest.skip("UVA cannot be used without GPUs.")
num_embs = 10 num_embs = 10
device=F.ctx() device = F.ctx()
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test_uva{}'.format(use_uva)) dgl_emb = NodeEmbedding(num_embs, emb_dim, "test_uva{}".format(use_uva))
torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(torch_emb.weight, 0, 1.0) th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
...@@ -69,7 +72,7 @@ def test_sparse_adam_uva(use_uva, emb_dim): ...@@ -69,7 +72,7 @@ def test_sparse_adam_uva(use_uva, emb_dim):
# first step # first step
idx = th.randint(0, num_embs, size=(4,)) idx = th.randint(0, num_embs, size=(4,))
dgl_value = dgl_emb(idx, device).to(th.device('cpu')) dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
torch_value = torch_emb(idx) torch_value = torch_emb(idx)
labels = th.zeros((4,)).long() labels = th.zeros((4,)).long()
...@@ -88,13 +91,14 @@ def test_sparse_adam_uva(use_uva, emb_dim): ...@@ -88,13 +91,14 @@ def test_sparse_adam_uva(use_uva, emb_dim):
# Pytorch sparseAdam maintains a global step # Pytorch sparseAdam maintains a global step
# DGL sparseAdam use a per embedding step # DGL sparseAdam use a per embedding step
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@pytest.mark.parametrize('dtype', [th.float32, th.float16]) @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@pytest.mark.parametrize('emb_dim', [1, 4, 101, 1024]) @pytest.mark.parametrize("dtype", [th.float32, th.float16])
@pytest.mark.parametrize("emb_dim", [1, 4, 101, 1024])
def test_sparse_adam_dtype(dtype, emb_dim): def test_sparse_adam_dtype(dtype, emb_dim):
num_embs = 10 num_embs = 10
device=F.ctx() device = F.ctx()
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test_dtype{}'.format(dtype)) dgl_emb = NodeEmbedding(num_embs, emb_dim, "test_dtype{}".format(dtype))
torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(torch_emb.weight, 0, 1.0) th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
...@@ -106,7 +110,7 @@ def test_sparse_adam_dtype(dtype, emb_dim): ...@@ -106,7 +110,7 @@ def test_sparse_adam_dtype(dtype, emb_dim):
# first step # first step
idx = th.randint(0, num_embs, size=(4,)) idx = th.randint(0, num_embs, size=(4,))
dgl_value = dgl_emb(idx, device).to(th.device('cpu')) dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
torch_value = torch_emb(idx) torch_value = torch_emb(idx)
labels = th.zeros((4,)).long() labels = th.zeros((4,)).long()
...@@ -126,15 +130,14 @@ def test_sparse_adam_dtype(dtype, emb_dim): ...@@ -126,15 +130,14 @@ def test_sparse_adam_dtype(dtype, emb_dim):
# DGL sparseAdam use a per embedding step # DGL sparseAdam use a per embedding step
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
def test_sparse_adam_zero_step(): def test_sparse_adam_zero_step():
num_embs = 10 num_embs = 10
emb_dim = 4 emb_dim = 4
device=F.ctx() device = F.ctx()
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test') dgl_emb = NodeEmbedding(num_embs, emb_dim, "test")
torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
dgl_emb_zero = NodeEmbedding(num_embs, emb_dim, 'test2') dgl_emb_zero = NodeEmbedding(num_embs, emb_dim, "test2")
torch_emb_zero = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb_zero = th.nn.Embedding(num_embs, emb_dim, sparse=True)
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(torch_emb.weight, 0, 1.0) th.nn.init.uniform_(torch_emb.weight, 0, 1.0)
...@@ -145,11 +148,13 @@ def test_sparse_adam_zero_step(): ...@@ -145,11 +148,13 @@ def test_sparse_adam_zero_step():
dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01) dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
torch_adam = th.optim.SparseAdam( torch_adam = th.optim.SparseAdam(
list(torch_emb.parameters()) + list(torch_emb_zero.parameters()), lr=0.01) list(torch_emb.parameters()) + list(torch_emb_zero.parameters()),
lr=0.01,
)
# first step # first step
idx = th.randint(0, num_embs, size=(4,)) idx = th.randint(0, num_embs, size=(4,))
dgl_value = dgl_emb(idx, device).to(th.device('cpu')) dgl_value = dgl_emb(idx, device).to(th.device("cpu"))
torch_value = torch_emb(idx) torch_value = torch_emb(idx)
labels = th.ones((4,)).long() labels = th.ones((4,)).long()
...@@ -164,33 +169,51 @@ def test_sparse_adam_zero_step(): ...@@ -164,33 +169,51 @@ def test_sparse_adam_zero_step():
torch_adam.step() torch_adam.step()
assert F.allclose(dgl_emb.weight, torch_emb.weight) assert F.allclose(dgl_emb.weight, torch_emb.weight)
def initializer(emb): def initializer(emb):
th.manual_seed(0) th.manual_seed(0)
emb.uniform_(-1.0, 1.0) emb.uniform_(-1.0, 1.0)
return emb return emb
def start_sparse_adam_worker(rank, device, world_size, weight, tensor_dev='cpu', has_zero_grad=False,
backend='gloo', num_embs=128, emb_dim=10):
print('start sparse worker for adam {}'.format(rank))
dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
master_ip='127.0.0.1', master_port='12345')
if device.type == 'cuda': def start_sparse_adam_worker(
rank,
device,
world_size,
weight,
tensor_dev="cpu",
has_zero_grad=False,
backend="gloo",
num_embs=128,
emb_dim=10,
):
print("start sparse worker for adam {}".format(rank))
dist_init_method = "tcp://{master_ip}:{master_port}".format(
master_ip="127.0.0.1", master_port="12345"
)
if device.type == "cuda":
th.cuda.set_device(device) th.cuda.set_device(device)
th.distributed.init_process_group(backend=backend, th.distributed.init_process_group(
init_method=dist_init_method, backend=backend,
world_size=world_size, init_method=dist_init_method,
rank=rank) world_size=world_size,
rank=rank,
)
init_weight = th.empty((num_embs, emb_dim)) init_weight = th.empty((num_embs, emb_dim))
th.manual_seed(0) th.manual_seed(0)
th.nn.init.uniform_(init_weight, -1.0, 1.0) th.nn.init.uniform_(init_weight, -1.0, 1.0)
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test', init_func=initializer, device=tensor_dev) dgl_emb = NodeEmbedding(
num_embs, emb_dim, "test", init_func=initializer, device=tensor_dev
)
dgl_emb.all_set_embedding(init_weight) dgl_emb.all_set_embedding(init_weight)
if has_zero_grad: if has_zero_grad:
dgl_emb_zero = NodeEmbedding(num_embs, emb_dim, 'zero', init_func=initializer, device=tensor_dev) dgl_emb_zero = NodeEmbedding(
num_embs, emb_dim, "zero", init_func=initializer, device=tensor_dev
)
dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01) dgl_adam = SparseAdam(params=[dgl_emb, dgl_emb_zero], lr=0.01)
else: else:
dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01) dgl_adam = SparseAdam(params=[dgl_emb], lr=0.01)
...@@ -215,17 +238,22 @@ def start_sparse_adam_worker(rank, device, world_size, weight, tensor_dev='cpu', ...@@ -215,17 +238,22 @@ def start_sparse_adam_worker(rank, device, world_size, weight, tensor_dev='cpu',
weight[:] = dgl_weight[:] weight[:] = dgl_weight[:]
th.distributed.barrier() th.distributed.barrier()
def start_torch_adam_worker(rank, world_size, weight, has_zero_grad=False,
num_embs=128, emb_dim=10):
print('start sparse worker for adam {}'.format(rank))
dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
master_ip='127.0.0.1', master_port='12345')
backend='gloo'
th.distributed.init_process_group(backend=backend, def start_torch_adam_worker(
init_method=dist_init_method, rank, world_size, weight, has_zero_grad=False, num_embs=128, emb_dim=10
world_size=world_size, ):
rank=rank) print("start sparse worker for adam {}".format(rank))
dist_init_method = "tcp://{master_ip}:{master_port}".format(
master_ip="127.0.0.1", master_port="12345"
)
backend = "gloo"
th.distributed.init_process_group(
backend=backend,
init_method=dist_init_method,
world_size=world_size,
rank=rank,
)
torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True) torch_emb = th.nn.Embedding(num_embs, emb_dim, sparse=True)
th.manual_seed(0) th.manual_seed(0)
...@@ -238,10 +266,14 @@ def start_torch_adam_worker(rank, world_size, weight, has_zero_grad=False, ...@@ -238,10 +266,14 @@ def start_torch_adam_worker(rank, world_size, weight, has_zero_grad=False,
th.nn.init.uniform_(torch_emb_zero.weight, -1.0, 1.0) th.nn.init.uniform_(torch_emb_zero.weight, -1.0, 1.0)
torch_emb_zero = th.nn.parallel.DistributedDataParallel(torch_emb_zero) torch_emb_zero = th.nn.parallel.DistributedDataParallel(torch_emb_zero)
torch_adam = th.optim.SparseAdam( torch_adam = th.optim.SparseAdam(
list(torch_emb.module.parameters()) + list(torch_emb_zero.module.parameters()), list(torch_emb.module.parameters())
lr=0.01) + list(torch_emb_zero.module.parameters()),
lr=0.01,
)
else: else:
torch_adam = th.optim.SparseAdam(list(torch_emb.module.parameters()), lr=0.01) torch_adam = th.optim.SparseAdam(
list(torch_emb.module.parameters()), lr=0.01
)
start = (num_embs // world_size) * rank start = (num_embs // world_size) * rank
end = (num_embs // world_size) * (rank + 1) end = (num_embs // world_size) * (rank + 1)
...@@ -259,20 +291,31 @@ def start_torch_adam_worker(rank, world_size, weight, has_zero_grad=False, ...@@ -259,20 +291,31 @@ def start_torch_adam_worker(rank, world_size, weight, has_zero_grad=False,
weight[:] = torch_emb.module.weight.cpu()[:] weight[:] = torch_emb.module.weight.cpu()[:]
th.distributed.barrier() th.distributed.barrier()
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type != 'cpu', reason='cpu only test') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(F.ctx().type != "cpu", reason="cpu only test")
@pytest.mark.parametrize("num_workers", [2, 4]) @pytest.mark.parametrize("num_workers", [2, 4])
def test_multiprocess_cpu_sparse_adam(num_workers): def test_multiprocess_cpu_sparse_adam(num_workers):
backend = 'gloo' backend = "gloo"
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = F.ctx() device = F.ctx()
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, th.device('cpu'), True, backend)) target=start_sparse_adam_worker,
args=(
i,
device,
num_workers,
dgl_weight,
th.device("cpu"),
True,
backend,
),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -281,8 +324,10 @@ def test_multiprocess_cpu_sparse_adam(num_workers): ...@@ -281,8 +324,10 @@ def test_multiprocess_cpu_sparse_adam(num_workers):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -290,26 +335,37 @@ def test_multiprocess_cpu_sparse_adam(num_workers): ...@@ -290,26 +335,37 @@ def test_multiprocess_cpu_sparse_adam(num_workers):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(F.ctx().type == "cpu", reason="gpu only test")
@pytest.mark.parametrize("num_workers", [2, 4, 8]) @pytest.mark.parametrize("num_workers", [2, 4, 8])
@pytest.mark.parametrize("backend", ['nccl', 'gloo']) @pytest.mark.parametrize("backend", ["nccl", "gloo"])
def test_multiprocess_sparse_adam(num_workers, backend): def test_multiprocess_sparse_adam(num_workers, backend):
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers: if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
pytest.skip("Not enough GPUs to run test.") pytest.skip("Not enough GPUs to run test.")
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = F.ctx() device = F.ctx()
if device.type == 'cuda': if device.type == "cuda":
# make sure each process has a unique GPU # make sure each process has a unique GPU
device = th.device(i) device = th.device(i)
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, th.device('cpu'), True, backend)) target=start_sparse_adam_worker,
args=(
i,
device,
num_workers,
dgl_weight,
th.device("cpu"),
True,
backend,
),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -318,8 +374,10 @@ def test_multiprocess_sparse_adam(num_workers, backend): ...@@ -318,8 +374,10 @@ def test_multiprocess_sparse_adam(num_workers, backend):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -327,25 +385,30 @@ def test_multiprocess_sparse_adam(num_workers, backend): ...@@ -327,25 +385,30 @@ def test_multiprocess_sparse_adam(num_workers, backend):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type == 'cpu', reason='cuda tensor is not supported for cpu') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(
F.ctx().type == "cpu", reason="cuda tensor is not supported for cpu"
)
@pytest.mark.parametrize("num_workers", [2, 4, 8]) @pytest.mark.parametrize("num_workers", [2, 4, 8])
def test_multiprocess_sparse_adam_cuda_tensor(num_workers): def test_multiprocess_sparse_adam_cuda_tensor(num_workers):
if F.ctx().type == 'cpu': if F.ctx().type == "cpu":
pytest.skip("Do not test CPU") pytest.skip("Do not test CPU")
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers: if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
pytest.skip("Not enough GPUs to run test.") pytest.skip("Not enough GPUs to run test.")
backend = 'nccl' backend = "nccl"
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = th.device(i) device = th.device(i)
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, device, False, backend)) target=start_sparse_adam_worker,
args=(i, device, num_workers, dgl_weight, device, False, backend),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -354,8 +417,10 @@ def test_multiprocess_sparse_adam_cuda_tensor(num_workers): ...@@ -354,8 +417,10 @@ def test_multiprocess_sparse_adam_cuda_tensor(num_workers):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -363,21 +428,32 @@ def test_multiprocess_sparse_adam_cuda_tensor(num_workers): ...@@ -363,21 +428,32 @@ def test_multiprocess_sparse_adam_cuda_tensor(num_workers):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type != 'cpu', reason='cpu only test') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(F.ctx().type != "cpu", reason="cpu only test")
@pytest.mark.parametrize("num_workers", [2, 4]) @pytest.mark.parametrize("num_workers", [2, 4])
def test_multiprocess_sparse_adam_cpu_zero_step(num_workers): def test_multiprocess_sparse_adam_cpu_zero_step(num_workers):
backend = 'gloo' backend = "gloo"
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = F.ctx() device = F.ctx()
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, th.device('cpu'), True, backend)) target=start_sparse_adam_worker,
args=(
i,
device,
num_workers,
dgl_weight,
th.device("cpu"),
True,
backend,
),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -386,8 +462,10 @@ def test_multiprocess_sparse_adam_cpu_zero_step(num_workers): ...@@ -386,8 +462,10 @@ def test_multiprocess_sparse_adam_cpu_zero_step(num_workers):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -395,26 +473,37 @@ def test_multiprocess_sparse_adam_cpu_zero_step(num_workers): ...@@ -395,26 +473,37 @@ def test_multiprocess_sparse_adam_cpu_zero_step(num_workers):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(F.ctx().type == "cpu", reason="gpu only test")
@pytest.mark.parametrize("num_workers", [2, 4, 8]) @pytest.mark.parametrize("num_workers", [2, 4, 8])
@pytest.mark.parametrize("backend", ['nccl', 'gloo']) @pytest.mark.parametrize("backend", ["nccl", "gloo"])
def test_multiprocess_sparse_adam_zero_step(num_workers, backend): def test_multiprocess_sparse_adam_zero_step(num_workers, backend):
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers: if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
pytest.skip("Not enough GPUs to run test.") pytest.skip("Not enough GPUs to run test.")
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = F.ctx() device = F.ctx()
if device.type == 'cuda': if device.type == "cuda":
# make sure each process has a unique GPU # make sure each process has a unique GPU
device = th.device(i) device = th.device(i)
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, th.device('cpu'), True, backend)) target=start_sparse_adam_worker,
args=(
i,
device,
num_workers,
dgl_weight,
th.device("cpu"),
True,
backend,
),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -423,8 +512,10 @@ def test_multiprocess_sparse_adam_zero_step(num_workers, backend): ...@@ -423,8 +512,10 @@ def test_multiprocess_sparse_adam_zero_step(num_workers, backend):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -432,23 +523,28 @@ def test_multiprocess_sparse_adam_zero_step(num_workers, backend): ...@@ -432,23 +523,28 @@ def test_multiprocess_sparse_adam_zero_step(num_workers, backend):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type == 'cpu', reason='cuda tensor is not supported for cpu') @unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@unittest.skipIf(
F.ctx().type == "cpu", reason="cuda tensor is not supported for cpu"
)
@pytest.mark.parametrize("num_workers", [2, 4, 8]) @pytest.mark.parametrize("num_workers", [2, 4, 8])
def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers): def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers):
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers: if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
pytest.skip("Not enough GPUs to run test.") pytest.skip("Not enough GPUs to run test.")
backend = 'nccl' backend = "nccl"
worker_list = [] worker_list = []
num_embs=128 num_embs = 128
emb_dim=10 emb_dim = 10
dgl_weight = th.empty((num_embs, emb_dim)) dgl_weight = th.empty((num_embs, emb_dim))
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
device = th.device(i) device = th.device(i)
p = ctx.Process(target=start_sparse_adam_worker, p = ctx.Process(
args=(i, device, num_workers, dgl_weight, device, True, backend)) target=start_sparse_adam_worker,
args=(i, device, num_workers, dgl_weight, device, True, backend),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -457,8 +553,10 @@ def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers): ...@@ -457,8 +553,10 @@ def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers):
worker_list = [] worker_list = []
torch_weight = th.empty((num_embs, emb_dim)) torch_weight = th.empty((num_embs, emb_dim))
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_torch_adam_worker, p = ctx.Process(
args=(i, num_workers, torch_weight, False)) target=start_torch_adam_worker,
args=(i, num_workers, torch_weight, False),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
for p in worker_list: for p in worker_list:
...@@ -466,7 +564,8 @@ def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers): ...@@ -466,7 +564,8 @@ def test_multiprocess_sparse_adam_zero_step_cuda_tensor(num_workers):
assert F.allclose(dgl_weight, torch_weight) assert F.allclose(dgl_weight, torch_weight)
if __name__ == '__main__':
if __name__ == "__main__":
test_sparse_adam(1) test_sparse_adam(1)
test_sparse_adam(4) test_sparse_adam(4)
test_sparse_adam(101) test_sparse_adam(101)
...@@ -478,15 +577,15 @@ if __name__ == '__main__': ...@@ -478,15 +577,15 @@ if __name__ == '__main__':
test_multiprocess_cpu_sparse_adam(8) test_multiprocess_cpu_sparse_adam(8)
test_multiprocess_sparse_adam_cpu_zero_step(2) test_multiprocess_sparse_adam_cpu_zero_step(2)
test_multiprocess_sparse_adam(2, backend='gloo') test_multiprocess_sparse_adam(2, backend="gloo")
test_multiprocess_sparse_adam(4, backend='gloo') test_multiprocess_sparse_adam(4, backend="gloo")
test_multiprocess_sparse_adam(8, backend='gloo') test_multiprocess_sparse_adam(8, backend="gloo")
test_multiprocess_sparse_adam(2, backend='nccl') test_multiprocess_sparse_adam(2, backend="nccl")
test_multiprocess_sparse_adam(4, backend='nccl') test_multiprocess_sparse_adam(4, backend="nccl")
test_multiprocess_sparse_adam(8, backend='nccl') test_multiprocess_sparse_adam(8, backend="nccl")
test_multiprocess_sparse_adam_zero_step(2, backend='gloo') test_multiprocess_sparse_adam_zero_step(2, backend="gloo")
test_multiprocess_sparse_adam_zero_step(4, backend='nccl') test_multiprocess_sparse_adam_zero_step(4, backend="nccl")
test_multiprocess_sparse_adam_cuda_tensor(2) test_multiprocess_sparse_adam_cuda_tensor(2)
test_multiprocess_sparse_adam_zero_step_cuda_tensor(4) test_multiprocess_sparse_adam_zero_step_cuda_tensor(4)
tests/pytorch/test_pickle.py
View file @ 89a4cc4d
import io
import pickle
import networkx as nx import networkx as nx
import dgl
import torch import torch
import pickle
import io import dgl
def _reconstruct_pickle(obj): def _reconstruct_pickle(obj):
f = io.BytesIO() f = io.BytesIO()
...@@ -12,15 +15,17 @@ def _reconstruct_pickle(obj): ...@@ -12,15 +15,17 @@ def _reconstruct_pickle(obj):
f.close() f.close()
return obj return obj
def test_pickling_batched_graph(): def test_pickling_batched_graph():
# NOTE: this is a test for a wierd bug mentioned in # NOTE: this is a test for a wierd bug mentioned in
# https://github.com/dmlc/dgl/issues/438 # https://github.com/dmlc/dgl/issues/438
glist = [nx.path_graph(i + 5) for i in range(5)] glist = [nx.path_graph(i + 5) for i in range(5)]
glist = [dgl.DGLGraph(g) for g in glist] glist = [dgl.DGLGraph(g) for g in glist]
bg = dgl.batch(glist) bg = dgl.batch(glist)
bg.ndata['x'] = torch.randn((35, 5)) bg.ndata["x"] = torch.randn((35, 5))
bg.edata['y'] = torch.randn((60, 3)) bg.edata["y"] = torch.randn((60, 3))
new_bg = _reconstruct_pickle(bg) new_bg = _reconstruct_pickle(bg)
if __name__ == '__main__':
if __name__ == "__main__":
test_pickling_batched_graph() test_pickling_batched_graph()
tests/pytorch/test_pin_memory.py
View file @ 89a4cc4d
import backend as F import backend as F
import dgl
import pytest import pytest
import torch import torch
@pytest.mark.skipif(F._default_context_str == 'cpu', reason="Need gpu for this test.") import dgl
@pytest.mark.skipif(
F._default_context_str == "cpu", reason="Need gpu for this test."
)
def test_pin_noncontiguous(): def test_pin_noncontiguous():
t = torch.empty([10, 100]).transpose(0, 1) t = torch.empty([10, 100]).transpose(0, 1)
...@@ -13,7 +17,10 @@ def test_pin_noncontiguous(): ...@@ -13,7 +17,10 @@ def test_pin_noncontiguous():
with pytest.raises(dgl.DGLError): with pytest.raises(dgl.DGLError):
dgl.utils.pin_memory_inplace(t) dgl.utils.pin_memory_inplace(t)
@pytest.mark.skipif(F._default_context_str == 'cpu', reason="Need gpu for this test.")
@pytest.mark.skipif(
F._default_context_str == "cpu", reason="Need gpu for this test."
)
def test_pin_view(): def test_pin_view():
t = torch.empty([100, 10]) t = torch.empty([100, 10])
v = t[10:20] v = t[10:20]
...@@ -24,7 +31,10 @@ def test_pin_view(): ...@@ -24,7 +31,10 @@ def test_pin_view():
with pytest.raises(dgl.DGLError): with pytest.raises(dgl.DGLError):
dgl.utils.pin_memory_inplace(v) dgl.utils.pin_memory_inplace(v)
@pytest.mark.skipif(F._default_context_str == 'cpu', reason='Need gpu for this test.')
@pytest.mark.skipif(
F._default_context_str == "cpu", reason="Need gpu for this test."
)
def test_unpin_automatically(): def test_unpin_automatically():
# run a sufficient number of iterations such that the memory pool should be # run a sufficient number of iterations such that the memory pool should be
# re-used # re-used
...@@ -38,26 +48,32 @@ def test_unpin_automatically(): ...@@ -38,26 +48,32 @@ def test_unpin_automatically():
assert not F.is_pinned(t) assert not F.is_pinned(t)
del t del t
@pytest.mark.skipif(F._default_context_str == 'cpu', reason='Need gpu for this test.')
@pytest.mark.skipif(
F._default_context_str == "cpu", reason="Need gpu for this test."
)
def test_pin_unpin_column(): def test_pin_unpin_column():
g = dgl.graph(([1, 2, 3, 4], [0, 0, 0, 0])) g = dgl.graph(([1, 2, 3, 4], [0, 0, 0, 0]))
g.ndata['x'] = torch.randn(g.num_nodes()) g.ndata["x"] = torch.randn(g.num_nodes())
g.pin_memory_() g.pin_memory_()
assert g.is_pinned() assert g.is_pinned()
assert g.ndata['x'].is_pinned() assert g.ndata["x"].is_pinned()
for col in g._node_frames[0].values(): for col in g._node_frames[0].values():
assert col.pinned_by_dgl assert col.pinned_by_dgl
assert col._data_nd is not None assert col._data_nd is not None
g.ndata['x'] = torch.randn(g.num_nodes()) # unpin the old ndata['x'] g.ndata["x"] = torch.randn(g.num_nodes()) # unpin the old ndata['x']
assert g.is_pinned() assert g.is_pinned()
for col in g._node_frames[0].values(): for col in g._node_frames[0].values():
assert not col.pinned_by_dgl assert not col.pinned_by_dgl
assert col._data_nd is None assert col._data_nd is None
assert not g.ndata['x'].is_pinned() assert not g.ndata["x"].is_pinned()
@pytest.mark.skipif(F._default_context_str == 'cpu', reason='Need gpu for this test.')
@pytest.mark.skipif(
F._default_context_str == "cpu", reason="Need gpu for this test."
)
def test_pin_empty(): def test_pin_empty():
t = torch.tensor([]) t = torch.tensor([])
assert not t.is_pinned() assert not t.is_pinned()
...@@ -68,6 +84,7 @@ def test_pin_empty(): ...@@ -68,6 +84,7 @@ def test_pin_empty():
nd = dgl.utils.pin_memory_inplace(t) nd = dgl.utils.pin_memory_inplace(t)
assert not t.is_pinned() assert not t.is_pinned()
if __name__ == "__main__": if __name__ == "__main__":
test_pin_noncontiguous() test_pin_noncontiguous()
test_pin_view() test_pin_view()
......
tests/pytorch/test_sparse_emb.py
View file @ 89a4cc4d
import multiprocessing as mp import multiprocessing as mp
import unittest, os import os
import pytest import unittest
import torch as th
import backend as F import backend as F
import pytest
import torch as th
from dgl.nn import NodeEmbedding from dgl.nn import NodeEmbedding
...@@ -13,46 +14,54 @@ def initializer(emb): ...@@ -13,46 +14,54 @@ def initializer(emb):
emb.uniform_(-1.0, 1.0) emb.uniform_(-1.0, 1.0)
return emb return emb
def check_all_set_all_get_func(device, init_emb): def check_all_set_all_get_func(device, init_emb):
num_embs = init_emb.shape[0] num_embs = init_emb.shape[0]
emb_dim = init_emb.shape[1] emb_dim = init_emb.shape[1]
dgl_emb = NodeEmbedding(num_embs, emb_dim, 'test', device=device) dgl_emb = NodeEmbedding(num_embs, emb_dim, "test", device=device)
dgl_emb.all_set_embedding(init_emb) dgl_emb.all_set_embedding(init_emb)
out_emb = dgl_emb.all_get_embedding() out_emb = dgl_emb.all_get_embedding()
assert F.allclose(init_emb, out_emb) assert F.allclose(init_emb, out_emb)
def start_sparse_worker(rank, world_size, test, args): def start_sparse_worker(rank, world_size, test, args):
print('start sparse worker {}'.format(rank)) print("start sparse worker {}".format(rank))
dist_init_method = 'tcp://{master_ip}:{master_port}'.format( dist_init_method = "tcp://{master_ip}:{master_port}".format(
master_ip='127.0.0.1', master_port='12345') master_ip="127.0.0.1", master_port="12345"
backend = 'gloo' )
backend = "gloo"
device = F.ctx() device = F.ctx()
if device.type == 'cuda': if device.type == "cuda":
device = th.device(rank) device = th.device(rank)
th.cuda.set_device(device) th.cuda.set_device(device)
th.distributed.init_process_group(backend=backend, th.distributed.init_process_group(
init_method=dist_init_method, backend=backend,
world_size=world_size, init_method=dist_init_method,
rank=rank) world_size=world_size,
rank=rank,
)
test(device, *args) test(device, *args)
th.distributed.barrier() th.distributed.barrier()
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(os.name == "nt", reason="Do not support windows yet")
@pytest.mark.parametrize("num_workers", [1, 2, 3]) @pytest.mark.parametrize("num_workers", [1, 2, 3])
def test_multiprocess_sparse_emb_get_set(num_workers): def test_multiprocess_sparse_emb_get_set(num_workers):
if F.ctx().type == 'cuda' and th.cuda.device_count() < num_workers: if F.ctx().type == "cuda" and th.cuda.device_count() < num_workers:
pytest.skip("Not enough GPUs to run test.") pytest.skip("Not enough GPUs to run test.")
worker_list = [] worker_list = []
init_emb = th.rand([1000, 8]) init_emb = th.rand([1000, 8])
ctx = mp.get_context('spawn') ctx = mp.get_context("spawn")
for i in range(num_workers): for i in range(num_workers):
p = ctx.Process(target=start_sparse_worker, p = ctx.Process(
args=(i, num_workers, check_all_set_all_get_func, (init_emb,))) target=start_sparse_worker,
args=(i, num_workers, check_all_set_all_get_func, (init_emb,)),
)
p.start() p.start()
worker_list.append(p) worker_list.append(p)
...@@ -62,7 +71,7 @@ def test_multiprocess_sparse_emb_get_set(num_workers): ...@@ -62,7 +71,7 @@ def test_multiprocess_sparse_emb_get_set(num_workers):
assert p.exitcode == 0 assert p.exitcode == 0
if __name__ == '__main__': if __name__ == "__main__":
test_sparse_emb_get_set(1) test_sparse_emb_get_set(1)
test_sparse_emb_get_set(2) test_sparse_emb_get_set(2)
test_sparse_emb_get_set(3) test_sparse_emb_get_set(3)
tests/pytorch/test_stream.py
View file @ 89a4cc4d
from statistics import mean
import unittest import unittest
from statistics import mean
import backend as F
import numpy as np import numpy as np
import torch import torch
import dgl import dgl
import dgl.ndarray as nd import dgl.ndarray as nd
from dgl import rand_graph
import dgl.ops as OPS import dgl.ops as OPS
from dgl._ffi.streams import to_dgl_stream_handle, _dgl_get_stream from dgl import rand_graph
from dgl._ffi.streams import _dgl_get_stream, to_dgl_stream_handle
from dgl.utils import to_dgl_context from dgl.utils import to_dgl_context
import backend as F
# borrowed from PyTorch, torch/testing/_internal/common_utils.py # borrowed from PyTorch, torch/testing/_internal/common_utils.py
def _get_cycles_per_ms() -> float: def _get_cycles_per_ms() -> float:
"""Measure and return approximate number of cycles per millisecond for torch.cuda._sleep """Measure and return approximate number of cycles per millisecond for torch.cuda._sleep"""
"""
def measure() -> float: def measure() -> float:
start = torch.cuda.Event(enable_timing=True) start = torch.cuda.Event(enable_timing=True)
...@@ -36,7 +38,10 @@ def _get_cycles_per_ms() -> float: ...@@ -36,7 +38,10 @@ def _get_cycles_per_ms() -> float:
vals = sorted(vals) vals = sorted(vals)
return mean(vals[2 : num - 2]) return mean(vals[2 : num - 2])
@unittest.skipIf(F._default_context_str == 'cpu', reason="stream only runs on GPU.")
@unittest.skipIf(
F._default_context_str == "cpu", reason="stream only runs on GPU."
)
def test_basics(): def test_basics():
g = rand_graph(10, 20, device=F.cpu()) g = rand_graph(10, 20, device=F.cpu())
x = torch.ones(g.num_nodes(), 10) x = torch.ones(g.num_nodes(), 10)
...@@ -57,22 +62,31 @@ def test_basics(): ...@@ -57,22 +62,31 @@ def test_basics():
s.synchronize() s.synchronize()
assert torch.equal(OPS.copy_u_sum(gg, xx), result) assert torch.equal(OPS.copy_u_sum(gg, xx), result)
@unittest.skipIf(F._default_context_str == 'cpu', reason="stream only runs on GPU.")
@unittest.skipIf(
F._default_context_str == "cpu", reason="stream only runs on GPU."
)
def test_set_get_stream(): def test_set_get_stream():
current_stream = torch.cuda.current_stream() current_stream = torch.cuda.current_stream()
# test setting another stream # test setting another stream
s = torch.cuda.Stream(device=F.ctx()) s = torch.cuda.Stream(device=F.ctx())
torch.cuda.set_stream(s) torch.cuda.set_stream(s)
assert to_dgl_stream_handle(s).value == _dgl_get_stream(to_dgl_context(F.ctx())).value assert (
to_dgl_stream_handle(s).value
== _dgl_get_stream(to_dgl_context(F.ctx())).value
)
# revert to default stream # revert to default stream
torch.cuda.set_stream(current_stream) torch.cuda.set_stream(current_stream)
@unittest.skipIf(F._default_context_str == 'cpu', reason="stream only runs on GPU.")
@unittest.skipIf(
F._default_context_str == "cpu", reason="stream only runs on GPU."
)
# borrowed from PyTorch, test/test_cuda.py: test_record_stream() # borrowed from PyTorch, test/test_cuda.py: test_record_stream()
def test_record_stream_ndarray(): def test_record_stream_ndarray():
cycles_per_ms = _get_cycles_per_ms() cycles_per_ms = _get_cycles_per_ms()
t = nd.array(np.array([1., 2., 3., 4.], dtype=np.float32), ctx=nd.cpu()) t = nd.array(np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), ctx=nd.cpu())
t.pin_memory_() t.pin_memory_()
result = nd.empty([4], ctx=nd.gpu(0)) result = nd.empty([4], ctx=nd.gpu(0))
stream = torch.cuda.Stream() stream = torch.cuda.Stream()
...@@ -84,25 +98,33 @@ def test_record_stream_ndarray(): ...@@ -84,25 +98,33 @@ def test_record_stream_ndarray():
tmp = t.copyto(nd.gpu(0)) tmp = t.copyto(nd.gpu(0))
ptr[0] = F.from_dgl_nd(tmp).data_ptr() ptr[0] = F.from_dgl_nd(tmp).data_ptr()
torch.cuda.current_stream().wait_stream(stream) torch.cuda.current_stream().wait_stream(stream)
tmp.record_stream( tmp.record_stream(to_dgl_stream_handle(torch.cuda.current_stream()))
to_dgl_stream_handle(torch.cuda.current_stream()))
torch.cuda._sleep(int(50 * cycles_per_ms)) # delay the copy torch.cuda._sleep(int(50 * cycles_per_ms)) # delay the copy
result.copyfrom(tmp) result.copyfrom(tmp)
perform_copy() perform_copy()
with torch.cuda.stream(stream): with torch.cuda.stream(stream):
tmp2 = nd.empty([4], ctx=nd.gpu(0)) tmp2 = nd.empty([4], ctx=nd.gpu(0))
assert F.from_dgl_nd(tmp2).data_ptr() != ptr[0], 'allocation re-used too soon' assert (
F.from_dgl_nd(tmp2).data_ptr() != ptr[0]
), "allocation re-used too soon"
assert torch.equal(F.from_dgl_nd(result).cpu(), torch.tensor([1., 2., 3., 4.])) assert torch.equal(
F.from_dgl_nd(result).cpu(), torch.tensor([1.0, 2.0, 3.0, 4.0])
)
# Check that the block will be re-used after the main stream finishes # Check that the block will be re-used after the main stream finishes
torch.cuda.current_stream().synchronize() torch.cuda.current_stream().synchronize()
with torch.cuda.stream(stream): with torch.cuda.stream(stream):
tmp3 = nd.empty([4], ctx=nd.gpu(0)) tmp3 = nd.empty([4], ctx=nd.gpu(0))
assert F.from_dgl_nd(tmp3).data_ptr() == ptr[0], 'allocation not re-used' assert (
F.from_dgl_nd(tmp3).data_ptr() == ptr[0]
), "allocation not re-used"
@unittest.skipIf(F._default_context_str == 'cpu', reason="stream only runs on GPU.")
@unittest.skipIf(
F._default_context_str == "cpu", reason="stream only runs on GPU."
)
def test_record_stream_graph_positive(): def test_record_stream_graph_positive():
cycles_per_ms = _get_cycles_per_ms() cycles_per_ms = _get_cycles_per_ms()
...@@ -133,7 +155,10 @@ def test_record_stream_graph_positive(): ...@@ -133,7 +155,10 @@ def test_record_stream_graph_positive():
torch.cuda.current_stream().synchronize() torch.cuda.current_stream().synchronize()
assert torch.equal(result, results2) assert torch.equal(result, results2)
@unittest.skipIf(F._default_context_str == 'cpu', reason="stream only runs on GPU.")
@unittest.skipIf(
F._default_context_str == "cpu", reason="stream only runs on GPU."
)
def test_record_stream_graph_negative(): def test_record_stream_graph_negative():
cycles_per_ms = _get_cycles_per_ms() cycles_per_ms = _get_cycles_per_ms()
...@@ -165,7 +190,8 @@ def test_record_stream_graph_negative(): ...@@ -165,7 +190,8 @@ def test_record_stream_graph_negative():
torch.cuda.current_stream().synchronize() torch.cuda.current_stream().synchronize()
assert not torch.equal(result, results2) assert not torch.equal(result, results2)
if __name__ == '__main__':
if __name__ == "__main__":
test_basics() test_basics()
test_set_get_stream() test_set_get_stream()
test_record_stream_ndarray() test_record_stream_ndarray()
......
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