Skip to content

GitLab

Unverified Commit e9b624fe authored by Minjie Wang's avatar Minjie Wang Committed by GitHub
Browse files

Merge branch 'master' into dist_part

parents 8086d1ed a88e7f7e
Hide whitespace changes
Inline Side-by-side
Showing with 831 additions and 206 deletions
tensoradapter/include/tensoradapter.h
View file @ e9b624fe
......@@ -18,7 +18,7 @@ namespace tensoradapter {
extern "C" {
/*!
* \brief Allocate an empty tensor
* \brief Allocate an empty tensor.
*
* \param shape The shape
* \param dtype The data type
......@@ -28,6 +28,24 @@ extern "C" {
DLManagedTensor* TAempty(
std::vector<int64_t> shape, DLDataType dtype, DLContext ctx);
#ifdef DGL_USE_CUDA
/*!
* \brief Allocate a piece of GPU memory via
* PyTorch's THCCachingAllocator.
*
* \param nbytes The size to be allocated.
* \return Pointer to the allocated memory.
*/
void* RawAlloc(size_t nbytes);
/*!
* \brief Free the GPU memory.
*
* \param ptr Pointer to the memory to be freed.
*/
void RawDelete(void* ptr);
#endif // DGL_USE_CUDA
}
}; // namespace tensoradapter
......
tensoradapter/pytorch/CMakeLists.txt
View file @ e9b624fe
......@@ -17,6 +17,10 @@ list(GET TORCH_PREFIX_VER 0 TORCH_PREFIX)
list(GET TORCH_PREFIX_VER 1 TORCH_VER)
message(STATUS "Configuring for PyTorch ${TORCH_VER}")
if(USE_CUDA)
add_definitions(-DDGL_USE_CUDA)
endif()
set(Torch_DIR "${TORCH_PREFIX}/Torch")
message(STATUS "Setting directory to ${Torch_DIR}")
find_package(Torch REQUIRED)
......
tensoradapter/pytorch/build.bat
View file @ e9b624fe
......@@ -11,7 +11,7 @@ IF x%1x == xx GOTO single
FOR %%X IN (%*) DO (
DEL /S /Q *
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
)
......@@ -21,7 +21,7 @@ GOTO end
:single
DEL /S /Q *
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% .. -G "Visual Studio 16 2019" || EXIT /B 1
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% .. -G "Visual Studio 16 2019" || EXIT /B 1
msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
......
tensoradapter/pytorch/build.sh
View file @ e9b624fe
......@@ -13,7 +13,7 @@ else
CPSOURCE=*.so
fi
CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST"
CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST -DUSE_CUDA=$USE_CUDA"
if [ $# -eq 0 ]; then
$CMAKE_COMMAND $CMAKE_FLAGS ..
......
tensoradapter/pytorch/torch.cpp
View file @ e9b624fe
......@@ -7,6 +7,9 @@
#include <tensoradapter_exports.h>
#include <torch/torch.h>
#include <ATen/DLConvertor.h>
#ifdef DGL_USE_CUDA
#include <c10/cuda/CUDACachingAllocator.h>
#endif // DGL_USE_CUDA
#include <vector>
#include <iostream>
......@@ -47,6 +50,16 @@ TA_EXPORTS DLManagedTensor* TAempty(
return at::toDLPack(tensor);
}
#ifdef DGL_USE_CUDA
TA_EXPORTS void* RawAlloc(size_t nbytes) {
return c10::cuda::CUDACachingAllocator::raw_alloc(nbytes);
}
TA_EXPORTS void RawDelete(void* ptr) {
c10::cuda::CUDACachingAllocator::raw_delete(ptr);
}
#endif // DGL_USE_CUDA
};
}; // namespace tensoradapter
tests/compute/test_csrmm.py
View file @ e9b624fe
......@@ -5,6 +5,10 @@ import dgl
from test_utils import parametrize_idtype
import backend as F
if F.backend_name == 'pytorch':
import torch
torch.backends.cuda.matmul.allow_tf32 = False
def _random_simple_graph(idtype, dtype, ctx, M, N, max_nnz, srctype, dsttype, etype):
src = np.random.randint(0, M, (max_nnz,))
dst = np.random.randint(0, N, (max_nnz,))
......
tests/compute/test_edge_softmax_hetero.py
View file @ e9b624fe
......@@ -2,6 +2,7 @@ import dgl
from dgl.ops import edge_softmax
import dgl.function as fn
from collections import Counter
import math
import numpy as np
import scipy.sparse as ssp
import itertools
......@@ -17,8 +18,6 @@ rfuncs = {'sum': fn.sum, 'max': fn.max, 'min': fn.min, 'mean': fn.mean}
fill_value = {'sum': 0, 'max': float("-inf")}
feat_size = 2
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
def create_test_heterograph(idtype):
# test heterograph from the docstring, plus a user -- wishes -- game relation
# 3 users, 2 games, 2 developers
......@@ -37,8 +36,26 @@ def create_test_heterograph(idtype):
assert g.idtype == idtype
assert g.device == F.ctx()
return g
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
def test_edge_softmax_unidirectional():
g = dgl.heterograph({
('A', 'AB', 'B'): ([1,2,3,1,2,3,1,2,3],[0,0,0,1,1,1,2,2,2]),
('B', 'BB', 'B'): ([0,1,2,0,1,2,0,1,2], [0,0,0,1,1,1,2,2,2])})
g = g.to(F.ctx())
g.edges['AB'].data['x'] = F.ones(9) * 2
g.edges['BB'].data['x'] = F.ones(9)
result = dgl.ops.edge_softmax(g, {'AB': g.edges['AB'].data['x'], 'BB': g.edges['BB'].data['x']})
ab = result['A', 'AB', 'B']
bb = result['B', 'BB', 'B']
e2 = F.zeros_like(ab) + math.exp(2) / ((math.exp(2) + math.exp(1)) * 3)
e1 = F.zeros_like(bb) + math.exp(1) / ((math.exp(2) + math.exp(1)) * 3)
assert F.allclose(ab, e2)
assert F.allclose(bb, e1)
@unittest.skipIf(dgl.backend.backend_name != 'pytorch', reason='Only support PyTorch for now')
@pytest.mark.parametrize('g', get_cases(['clique']))
@pytest.mark.parametrize('norm_by', ['src', 'dst'])
# @pytest.mark.parametrize('shp', edge_softmax_shapes)
......@@ -109,5 +126,4 @@ def test_edge_softmax(g, norm_by, idtype):
assert F.allclose(grad_edata_hm, grad_edata_ht)
if __name__ == '__main__':
test_edge_softmax()
test_edge_softmax_unidirectional()
tests/compute/test_sampling.py
View file @ e9b624fe
......@@ -24,59 +24,74 @@ def check_random_walk(g, metapath, traces, ntypes, prob=None, trace_eids=None):
u, v = g.find_edges(trace_eids[i, j], etype=metapath[j])
assert (u == traces[i, j]) and (v == traces[i, j + 1])
@unittest.skipIf(F._default_context_str == 'gpu', reason="Random walk with non-uniform prob is not supported in GPU.")
def test_non_uniform_random_walk():
@pytest.mark.parametrize('use_uva', [True, False])
def test_non_uniform_random_walk(use_uva):
if use_uva:
if F.ctx() == F.cpu():
pytest.skip('UVA biased random walk requires a GPU.')
if dgl.backend.backend_name != 'pytorch':
pytest.skip('UVA biased random walk is only supported with PyTorch.')
g2 = dgl.heterograph({
('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0])
}).to(F.ctx())
})
g4 = dgl.heterograph({
('user', 'follow', 'user'): ([0, 1, 1, 2, 3], [1, 2, 3, 0, 0]),
('user', 'view', 'item'): ([0, 0, 1, 2, 3, 3], [0, 1, 1, 2, 2, 1]),
('item', 'viewed-by', 'user'): ([0, 1, 1, 2, 2, 1], [0, 0, 1, 2, 3, 3])
}).to(F.ctx())
})
g2.edata['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g2.edata['p2'] = F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32)
g4.edges['follow'].data['p'] = F.tensor([3, 0, 3, 3, 3], dtype=F.float32)
g4.edges['viewed-by'].data['p'] = F.tensor([1, 1, 1, 1, 1, 1], dtype=F.float32)
g2.edata['p'] = F.copy_to(F.tensor([3, 0, 3, 3, 3], dtype=F.float32), F.cpu())
g2.edata['p2'] = F.copy_to(F.tensor([[3], [0], [3], [3], [3]], dtype=F.float32), F.cpu())
g4.edges['follow'].data['p'] = F.copy_to(F.tensor([3, 0, 3, 3, 3], dtype=F.float32), F.cpu())
g4.edges['viewed-by'].data['p'] = F.copy_to(F.tensor([1, 1, 1, 1, 1, 1], dtype=F.float32), F.cpu())
traces, eids, ntypes = dgl.sampling.random_walk(
g2, [0, 1, 2, 3, 0, 1, 2, 3], length=4, prob='p', return_eids=True)
check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p', trace_eids=eids)
if use_uva:
for g in (g2, g4):
g.create_formats_()
g.pin_memory_()
elif F._default_context_str == 'gpu':
g2 = g2.to(F.ctx())
g4 = g4.to(F.ctx())
try:
traces, ntypes = dgl.sampling.random_walk(
g2, [0, 1, 2, 3, 0, 1, 2, 3], length=4, prob='p2')
fail = False
except dgl.DGLError:
fail = True
assert fail
traces, eids, ntypes = dgl.sampling.random_walk(
g2, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g2.idtype),
length=4, prob='p', return_eids=True)
check_random_walk(g2, ['follow'] * 4, traces, ntypes, 'p', trace_eids=eids)
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, eids, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p', return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p', restart_prob=0., return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath, prob='p',
restart_prob=F.zeros((6,), F.float32, F.cpu()), return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, [0, 1, 2, 3, 0, 1, 2, 3], metapath=metapath + ['follow'], prob='p',
restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32), return_eids=True)
check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p', trace_eids=eids)
assert (F.asnumpy(traces[:, 7]) == -1).all()
def _use_uva():
if F._default_context_str == 'cpu':
return [False]
else:
return [True, False]
with pytest.raises(dgl.DGLError):
traces, ntypes = dgl.sampling.random_walk(
g2, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g2.idtype),
length=4, prob='p2')
@pytest.mark.parametrize('use_uva', _use_uva())
metapath = ['follow', 'view', 'viewed-by'] * 2
traces, eids, ntypes = dgl.sampling.random_walk(
g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
metapath=metapath, prob='p', return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
metapath=metapath, prob='p', restart_prob=0., return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
metapath=metapath, prob='p',
restart_prob=F.zeros((6,), F.float32, F.ctx()), return_eids=True)
check_random_walk(g4, metapath, traces, ntypes, 'p', trace_eids=eids)
traces, eids, ntypes = dgl.sampling.random_walk(
g4, F.tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype=g4.idtype),
metapath=metapath + ['follow'], prob='p',
restart_prob=F.tensor([0, 0, 0, 0, 0, 0, 1], F.float32), return_eids=True)
check_random_walk(g4, metapath, traces[:, :7], ntypes[:7], 'p', trace_eids=eids)
assert (F.asnumpy(traces[:, 7]) == -1).all()
finally:
for g in (g2, g4):
g.unpin_memory_()
@pytest.mark.parametrize('use_uva', [True, False])
def test_uniform_random_walk(use_uva):
if use_uva and F.ctx() == F.cpu():
pytest.skip('UVA random walk requires a GPU.')
g1 = dgl.heterograph({
('user', 'follow', 'user'): ([0, 1, 2], [1, 2, 0])
})
......@@ -179,8 +194,10 @@ def test_pack_traces():
assert F.array_equal(result[2], F.tensor([2, 7], dtype=F.int64))
assert F.array_equal(result[3], F.tensor([0, 2], dtype=F.int64))
@pytest.mark.parametrize('use_uva', _use_uva())
@pytest.mark.parametrize('use_uva', [True, False])
def test_pinsage_sampling(use_uva):
if use_uva and F.ctx() == F.cpu():
pytest.skip('UVA sampling requires a GPU.')
def _test_sampler(g, sampler, ntype):
seeds = F.copy_to(F.tensor([0, 2], dtype=g.idtype), F.ctx())
neighbor_g = sampler(seeds)
......@@ -626,12 +643,10 @@ def test_sample_neighbors_noprob():
_test_sample_neighbors(False, None)
#_test_sample_neighbors(True)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors with probability is not implemented")
def test_sample_neighbors_prob():
_test_sample_neighbors(False, 'prob')
#_test_sample_neighbors(True)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
def test_sample_neighbors_outedge():
_test_sample_neighbors_outedge(False)
#_test_sample_neighbors_outedge(True)
......@@ -646,9 +661,8 @@ def test_sample_neighbors_topk_outedge():
_test_sample_neighbors_topk_outedge(False)
#_test_sample_neighbors_topk_outedge(True)
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
def test_sample_neighbors_with_0deg():
g = dgl.graph(([], []), num_nodes=5)
g = dgl.graph(([], []), num_nodes=5).to(F.ctx())
sg = dgl.sampling.sample_neighbors(g, F.tensor([1, 2], dtype=F.int64), 2, edge_dir='in', replace=False)
assert sg.number_of_edges() == 0
sg = dgl.sampling.sample_neighbors(g, F.tensor([1, 2], dtype=F.int64), 2, edge_dir='in', replace=True)
......@@ -885,7 +899,6 @@ def test_sample_neighbors_etype_sorted_homogeneous(format_, direction):
assert fail
@pytest.mark.parametrize('dtype', ['int32', 'int64'])
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
def test_sample_neighbors_exclude_edges_heteroG(dtype):
d_i_d_u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300, size=100, dtype=dtype)))
d_i_d_v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=d_i_d_u_nodes.shape, dtype=dtype))
......@@ -898,7 +911,7 @@ def test_sample_neighbors_exclude_edges_heteroG(dtype):
('drug', 'interacts', 'drug'): (d_i_d_u_nodes, d_i_d_v_nodes),
('drug', 'interacts', 'gene'): (d_i_g_u_nodes, d_i_g_v_nodes),
('drug', 'treats', 'disease'): (d_t_d_u_nodes, d_t_d_v_nodes)
})
}).to(F.ctx())
(U, V, EID) = (0, 1, 2)
......@@ -951,11 +964,10 @@ def test_sample_neighbors_exclude_edges_heteroG(dtype):
etype=('drug','treats','disease'))))
@pytest.mark.parametrize('dtype', ['int32', 'int64'])
@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU sample neighbors not implemented")
def test_sample_neighbors_exclude_edges_homoG(dtype):
u_nodes = F.zerocopy_from_numpy(np.unique(np.random.randint(300,size=100, dtype=dtype)))
v_nodes = F.zerocopy_from_numpy(np.random.randint(25, size=u_nodes.shape, dtype=dtype))
g = dgl.graph((u_nodes, v_nodes))
g = dgl.graph((u_nodes, v_nodes)).to(F.ctx())
(U, V, EID) = (0, 1, 2)
......
tests/compute/test_sparse.py
View file @ e9b624fe
......@@ -293,6 +293,8 @@ def test_segment_reduce(reducer):
@pytest.mark.parametrize('feat_size', [1, 8, 16, 64, 256])
@pytest.mark.parametrize('dtype,tol', [(torch.float16,1e-2),(torch.float32,3e-3),(torch.float64,1e-4)])
def test_segment_mm(idtype, feat_size, dtype, tol):
if F._default_context_str == 'cpu' and dtype == torch.float16:
pytest.skip("fp16 support for CPU linalg functions has been removed in PyTorch.")
dev = F.ctx()
# input
a = torch.tensor(np.random.rand(100, feat_size)).to(dev).to(dtype)
......
tests/compute/test_transform.py
View file @ e9b624fe
......@@ -1625,8 +1625,11 @@ def test_remove_nodes(idtype):
@parametrize_idtype
def test_add_selfloop(idtype):
# homogeneous graph
# test for fill_data is float
g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]), ctx=F.ctx())
g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.add_self_loop(g)
assert g.number_of_nodes() == 3
......@@ -1634,7 +1637,39 @@ def test_add_selfloop(idtype):
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 0, 0, 0], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
[1., 1.], [1., 1.], [1., 1.]]))
# test for fill_data is int
g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g.edata['he1'] = F.copy_to(F.tensor([[0, 1], [2, 3], [4, 5]], dtype=idtype), ctx=F.ctx())
g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.add_self_loop(g, fill_data=1)
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 6
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([1, 2, 3, 1, 1, 1], dtype=idtype))
assert F.array_equal(g.edata['he1'], F.tensor([[0, 1], [2, 3], [4, 5],
[1, 1], [1, 1], [1, 1]], dtype=idtype))
# test for fill_data is str
g = dgl.graph(([0, 0, 2], [2, 1, 0]), idtype=idtype, device=F.ctx())
g.edata['he'] = F.copy_to(F.tensor([1., 2., 3.]), ctx=F.ctx())
g.edata['he1'] = F.copy_to(F.tensor([[0., 1.], [2., 3.], [4., 5.]]), ctx=F.ctx())
g.ndata['hn'] = F.copy_to(F.tensor([1, 2, 3], dtype=idtype), ctx=F.ctx())
g = dgl.add_self_loop(g, fill_data='sum')
assert g.number_of_nodes() == 3
assert g.number_of_edges() == 6
u, v = g.edges(form='uv', order='eid')
assert F.array_equal(u, F.tensor([0, 0, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([2, 1, 0, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edata['he'], F.tensor([1., 2., 3., 3., 2., 1.]))
assert F.array_equal(g.edata['he1'], F.tensor([[0., 1.], [2., 3.], [4., 5.],
[4., 5.], [2., 3.], [0., 1.]]))
# bipartite graph
g = dgl.heterograph(
......@@ -1647,7 +1682,9 @@ def test_add_selfloop(idtype):
raise_error = True
assert raise_error
# test for fill_data is float
g = create_test_heterograph5(idtype)
g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]), ctx=F.ctx())
g = dgl.add_self_loop(g, etype='follows')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 2
......@@ -1656,9 +1693,52 @@ def test_add_selfloop(idtype):
u, v = g.edges(form='uv', order='eid', etype='follows')
assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 0, 0, 0], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 1, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0., 1.], [1., 2.], [1., 1.],
[1., 1.], [1., 1.]]))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
# test for fill_data is int
g = create_test_heterograph5(idtype)
g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0, 1], [1, 2]], dtype=idtype), ctx=F.ctx())
g = dgl.add_self_loop(g, fill_data=1, etype='follows')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 2
assert g.number_of_edges('follows') == 5
assert g.number_of_edges('plays') == 2
u, v = g.edges(form='uv', order='eid', etype='follows')
assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1, 2, 1, 1, 1], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0, 1], [1, 2], [1, 1],
[1, 1], [1, 1]], dtype=idtype))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1, 2], dtype=idtype))
# test for fill_data is str
g = dgl.heterograph({
('user', 'follows', 'user'): (F.tensor([1, 2], dtype=idtype),
F.tensor([0, 1], dtype=idtype)),
('user', 'plays', 'game'): (F.tensor([0, 1], dtype=idtype),
F.tensor([0, 1], dtype=idtype))},
idtype=idtype, device=F.ctx())
g.nodes['user'].data['h'] = F.copy_to(F.tensor([1, 1, 1], dtype=idtype), ctx=F.ctx())
g.nodes['game'].data['h'] = F.copy_to(F.tensor([2, 2], dtype=idtype), ctx=F.ctx())
g.edges['follows'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
g.edges['follows'].data['h1'] = F.copy_to(F.tensor([[0., 1.], [1., 2.]]), ctx=F.ctx())
g.edges['plays'].data['h'] = F.copy_to(F.tensor([1., 2.]), ctx=F.ctx())
g = dgl.add_self_loop(g, fill_data='mean', etype='follows')
assert g.number_of_nodes('user') == 3
assert g.number_of_nodes('game') == 2
assert g.number_of_edges('follows') == 5
assert g.number_of_edges('plays') == 2
u, v = g.edges(form='uv', order='eid', etype='follows')
assert F.array_equal(u, F.tensor([1, 2, 0, 1, 2], dtype=idtype))
assert F.array_equal(v, F.tensor([0, 1, 0, 1, 2], dtype=idtype))
assert F.array_equal(g.edges['follows'].data['h'], F.tensor([1., 2., 1., 2., 0.]))
assert F.array_equal(g.edges['follows'].data['h1'], F.tensor([[0., 1.], [1., 2.], [0., 1.],
[1., 2.], [0., 0.]]))
assert F.array_equal(g.edges['plays'].data['h'], F.tensor([1., 2.]))
raise_error = False
try:
g = dgl.add_self_loop(g, etype='plays')
......
tests/cpp/socket_communicator_test.cc
View file @ e9b624fe
......@@ -154,6 +154,14 @@ void start_server(int id) {
receiver.Finalize();
}
TEST(SocketCommunicatorTest, TCPSocketBind) {
dgl::network::TCPSocket socket;
testing::internal::CaptureStderr();
EXPECT_EQ(socket.Bind("127.0.0", 50001), false);
const std::string stderr = testing::internal::GetCapturedStderr();
EXPECT_NE(stderr.find("Invalid IP: 127.0.0"), std::string::npos);
}
#else
#include <windows.h>
......
tests/cugraph/test_basics.py 0 → 100644
View file @ e9b624fe
import backend as F
import dgl
import numpy as np
from dgl import DGLGraph
import unittest
import pytest
import cugraph
def test_dummy():
cg = cugraph.Graph()
assert cg is not None
def test_to_cugraph_conversion():
g = dgl.graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to('cuda')
cugraph_g = g.to_cugraph()
assert cugraph_g.number_of_nodes()==g.number_of_nodes()
assert cugraph_g.number_of_edges()==g.number_of_edges()
assert cugraph_g.has_edge(0, 1)
assert cugraph_g.has_edge(1, 0)
assert cugraph_g.has_edge(3, 2)
def test_from_cugraph_conversion():
# cudf is a dependency of cugraph
import cudf
# directed graph conversion test
cugraph_g = cugraph.Graph(directed=True)
df = cudf.DataFrame({"source":[0, 1, 2, 3],
"destination":[1, 2, 3, 2]})
cugraph_g.from_cudf_edgelist(df)
g = dgl.from_cugraph(cugraph_g)
assert g.device.type == 'cuda'
assert g.number_of_nodes() == cugraph_g.number_of_nodes()
assert g.number_of_edges() == cugraph_g.number_of_edges()
# assert reverse edges are not present
assert g.has_edges_between(0, 1)
assert not g.has_edges_between(1, 0)
assert g.has_edges_between(1, 2)
assert not g.has_edges_between(2, 1)
assert g.has_edges_between(2, 3)
# undirected graph conversion test
cugraph_g = cugraph.Graph(directed=False)
df = cudf.DataFrame({"source":[0, 1, 2, 3],
"destination":[1, 2, 3, 2]})
cugraph_g.from_cudf_edgelist(df)
g = dgl.from_cugraph(cugraph_g)
assert g.device.type == 'cuda'
assert g.number_of_nodes() == cugraph_g.number_of_nodes()
# assert reverse edges are present
assert g.has_edges_between(0, 1)
assert g.has_edges_between(1, 0)
assert g.has_edges_between(1, 2)
assert g.has_edges_between(2, 1)
assert g.has_edges_between(2, 3)
tests/dist/python/run_dist_objects.py 0 → 100644
View file @ e9b624fe
import dgl
import os
import numpy as np
import dgl.backend as F
from dgl.distributed import load_partition_book
mode = os.environ.get('DIST_DGL_TEST_MODE', "")
graph_name = os.environ.get('DIST_DGL_TEST_GRAPH_NAME', 'random_test_graph')
num_part = int(os.environ.get('DIST_DGL_TEST_NUM_PART'))
num_servers_per_machine = int(os.environ.get('DIST_DGL_TEST_NUM_SERVER'))
num_client_per_machine = int(os.environ.get('DIST_DGL_TEST_NUM_CLIENT'))
shared_workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE')
graph_path = os.environ.get('DIST_DGL_TEST_GRAPH_PATH')
part_id = int(os.environ.get('DIST_DGL_TEST_PART_ID'))
net_type = os.environ.get('DIST_DGL_TEST_NET_TYPE')
ip_config = os.environ.get('DIST_DGL_TEST_IP_CONFIG', 'ip_config.txt')
os.environ['DGL_DIST_MODE'] = 'distributed'
def zeros_init(shape, dtype):
return F.zeros(shape, dtype=dtype, ctx=F.cpu())
def run_server(graph_name, server_id, server_count, num_clients, shared_mem, keep_alive=False):
# server_count = num_servers_per_machine
g = dgl.distributed.DistGraphServer(server_id, ip_config,
server_count, num_clients,
graph_path + '/{}.json'.format(graph_name),
disable_shared_mem=not shared_mem,
graph_format=['csc', 'coo'], keep_alive=keep_alive,
net_type=net_type)
print('start server', server_id)
g.start()
##########################################
############### DistTensor ###############
##########################################
def dist_tensor_test_sanity(data_shape, name=None):
local_rank = dgl.distributed.get_rank() % num_client_per_machine
dist_ten = dgl.distributed.DistTensor(data_shape,
F.int32,
init_func=zeros_init,
name=name)
# arbitrary value
stride = 3
pos = (part_id // 2) * num_client_per_machine + local_rank
if part_id % 2 == 0:
dist_ten[pos*stride:(pos+1)*stride] = F.ones((stride, 2), dtype=F.int32, ctx=F.cpu()) * (pos+1)
dgl.distributed.client_barrier()
assert F.allclose(dist_ten[pos*stride:(pos+1)*stride],
F.ones((stride, 2), dtype=F.int32, ctx=F.cpu()) * (pos+1))
def dist_tensor_test_destroy_recreate(data_shape, name):
dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, name, init_func=zeros_init)
del dist_ten
dgl.distributed.client_barrier()
new_shape = (data_shape[0], 4)
dist_ten = dgl.distributed.DistTensor(new_shape, F.float32, name, init_func=zeros_init)
def dist_tensor_test_persistent(data_shape):
dist_ten_name = 'persistent_dist_tensor'
dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, dist_ten_name, init_func=zeros_init,
persistent=True)
del dist_ten
try:
dist_ten = dgl.distributed.DistTensor(data_shape, F.float32, dist_ten_name)
raise Exception('')
except:
pass
def test_dist_tensor(g):
first_type = g.ntypes[0]
data_shape = (g.number_of_nodes(first_type), 2)
dist_tensor_test_sanity(data_shape)
dist_tensor_test_sanity(data_shape, name="DistTensorSanity")
dist_tensor_test_destroy_recreate(data_shape, name="DistTensorRecreate")
dist_tensor_test_persistent(data_shape)
##########################################
############# DistEmbedding ##############
##########################################
def dist_embedding_check_sanity(num_nodes, optimizer, name=None):
local_rank = dgl.distributed.get_rank() % num_client_per_machine
emb = dgl.distributed.DistEmbedding(num_nodes, 1, name=name, init_func=zeros_init)
lr = 0.001
optim = optimizer(params=[emb], lr=lr)
stride = 3
pos = (part_id // 2) * num_client_per_machine + local_rank
idx = F.arange(pos*stride, (pos+1)*stride)
if part_id % 2 == 0:
with F.record_grad():
value = emb(idx)
optim.zero_grad()
loss = F.sum(value + 1, 0)
loss.backward()
optim.step()
dgl.distributed.client_barrier()
value = emb(idx)
F.allclose(value, F.ones((len(idx), 1), dtype=F.int32, ctx=F.cpu()) * -lr)
not_update_idx = F.arange(((num_part + 1) / 2) * num_client_per_machine * stride, num_nodes)
value = emb(not_update_idx)
assert np.all(F.asnumpy(value) == np.zeros((len(not_update_idx), 1)))
def dist_embedding_check_existing(num_nodes):
dist_emb_name = "UniqueEmb"
emb = dgl.distributed.DistEmbedding(num_nodes, 1, name=dist_emb_name, init_func=zeros_init)
try:
emb1 = dgl.distributed.DistEmbedding(num_nodes, 2, name=dist_emb_name, init_func=zeros_init)
raise Exception('')
except:
pass
def test_dist_embedding(g):
num_nodes = g.number_of_nodes(g.ntypes[0])
dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdagrad)
dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdagrad, name='SomeEmbedding')
dist_embedding_check_sanity(num_nodes, dgl.distributed.optim.SparseAdam, name='SomeEmbedding')
dist_embedding_check_existing(num_nodes)
if mode == "server":
shared_mem = bool(int(os.environ.get('DIST_DGL_TEST_SHARED_MEM')))
server_id = int(os.environ.get('DIST_DGL_TEST_SERVER_ID'))
run_server(graph_name, server_id, server_count=num_servers_per_machine,
num_clients=num_part*num_client_per_machine, shared_mem=shared_mem, keep_alive=False)
elif mode == "client":
os.environ['DGL_NUM_SERVER'] = str(num_servers_per_machine)
dgl.distributed.initialize(ip_config, net_type=net_type)
gpb, graph_name, _, _ = load_partition_book(graph_path + '/{}.json'.format(graph_name), part_id, None)
g = dgl.distributed.DistGraph(graph_name, gpb=gpb)
target_func_map = {"DistTensor": test_dist_tensor,
"DistEmbedding": test_dist_embedding,
}
target = os.environ.get("DIST_DGL_TEST_OBJECT_TYPE", "")
if target not in target_func_map:
for test_func in target_func_map.values():
test_func(g)
else:
target_func_map[target](g)
else:
print("DIST_DGL_TEST_MODE has to be either server or client")
exit(1)
tests/dist/test_dist_objects.py 0 → 100644
View file @ e9b624fe
import os
import unittest
import pytest
import multiprocessing as mp
import subprocess
import utils
import dgl
import numpy as np
import dgl.backend as F
from dgl.distributed import partition_graph
graph_name = os.environ.get('DIST_DGL_TEST_GRAPH_NAME', 'random_test_graph')
target = os.environ.get('DIST_DGL_TEST_OBJECT_TYPE', '')
shared_workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE')
def create_graph(num_part, dist_graph_path, hetero):
if not hetero:
g = dgl.rand_graph(10000, 42000)
g.ndata['feat'] = F.unsqueeze(F.arange(0, g.number_of_nodes()), 1)
g.edata['feat'] = F.unsqueeze(F.arange(0, g.number_of_edges()), 1)
partition_graph(g, graph_name, num_part, dist_graph_path)
else:
from scipy import sparse as spsp
num_nodes = {'n1': 10000, 'n2': 10010, 'n3': 10020}
etypes = [('n1', 'r1', 'n2'),
('n1', 'r2', 'n3'),
('n2', 'r3', 'n3')]
edges = {}
for etype in etypes:
src_ntype, _, dst_ntype = etype
arr = spsp.random(num_nodes[src_ntype], num_nodes[dst_ntype], density=0.001, format='coo',
random_state=100)
edges[etype] = (arr.row, arr.col)
g = dgl.heterograph(edges, num_nodes)
g.nodes['n1'].data['feat'] = F.unsqueeze(F.arange(0, g.number_of_nodes('n1')), 1)
g.edges['r1'].data['feat'] = F.unsqueeze(F.arange(0, g.number_of_edges('r1')), 1)
partition_graph(g, graph_name, num_part, dist_graph_path)
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@pytest.mark.parametrize("net_type", ['tensorpipe', 'socket'])
@pytest.mark.parametrize("num_servers", [1, 4])
@pytest.mark.parametrize("num_clients", [1, 4])
@pytest.mark.parametrize("hetero", [False, True])
@pytest.mark.parametrize("shared_mem", [False, True])
def test_dist_objects(net_type, num_servers, num_clients, hetero, shared_mem):
if not shared_mem and num_servers > 1:
pytest.skip(f"Backup servers are not supported when shared memory is disabled")
ip_config = os.environ.get('DIST_DGL_TEST_IP_CONFIG', 'ip_config.txt')
workspace = os.environ.get('DIST_DGL_TEST_WORKSPACE', '/shared_workspace/dgl_dist_tensor_test/')
ips = utils.get_ips(ip_config)
num_part = len(ips)
test_bin = os.path.join(os.environ.get(
'DIST_DGL_TEST_PY_BIN_DIR', '.'), 'run_dist_objects.py')
dist_graph_path = os.path.join(workspace, 'hetero_dist_graph' if hetero else 'dist_graph')
if not os.path.isdir(dist_graph_path):
create_graph(num_part, dist_graph_path, hetero)
base_envs = f"DIST_DGL_TEST_WORKSPACE={workspace} " \
f"DIST_DGL_TEST_NUM_PART={num_part} " \
f"DIST_DGL_TEST_NUM_SERVER={num_servers} " \
f"DIST_DGL_TEST_NUM_CLIENT={num_clients} " \
f"DIST_DGL_TEST_NET_TYPE={net_type} " \
f"DIST_DGL_TEST_GRAPH_PATH={dist_graph_path} " \
f"DIST_DGL_TEST_IP_CONFIG={ip_config} "
procs = []
# Start server
server_id = 0
for part_id, ip in enumerate(ips):
for _ in range(num_servers):
cmd_envs = base_envs + \
f"DIST_DGL_TEST_SERVER_ID={server_id} " \
f"DIST_DGL_TEST_PART_ID={part_id} " \
f"DIST_DGL_TEST_SHARED_MEM={str(int(shared_mem))} " \
f"DIST_DGL_TEST_MODE=server "
procs.append(utils.execute_remote(
f"{cmd_envs} python3 {test_bin}",
ip))
server_id += 1
# Start client processes
for part_id, ip in enumerate(ips):
for _ in range(num_clients):
cmd_envs = base_envs + \
f"DIST_DGL_TEST_PART_ID={part_id} " \
f"DIST_DGL_TEST_OBJECT_TYPE={target} " \
f"DIST_DGL_TEST_MODE=client "
procs.append(utils.execute_remote(
f"{cmd_envs} python3 {test_bin}",
ip))
for p in procs:
p.join()
assert p.exitcode == 0
tests/distributed/test_dist_graph_store.py
View file @ e9b624fe
......@@ -35,6 +35,15 @@ def run_server(graph_name, server_id, server_count, num_clients, shared_mem, kee
disable_shared_mem=not shared_mem,
graph_format=['csc', 'coo'], keep_alive=keep_alive)
print('start server', server_id)
# verify dtype of underlying graph
cg = g.client_g
for k, dtype in dgl.distributed.dist_graph.FIELD_DICT.items():
if k in cg.ndata:
assert F.dtype(
cg.ndata[k]) == dtype, "Data type of {} in ndata should be {}.".format(k, dtype)
if k in cg.edata:
assert F.dtype(
cg.edata[k]) == dtype, "Data type of {} in edata should be {}.".format(k, dtype)
g.start()
def emb_init(shape, dtype):
......@@ -219,6 +228,11 @@ def check_dist_graph(g, num_clients, num_nodes, num_edges):
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == eids))
# Test edge_subgraph
sg = g.edge_subgraph(eids)
assert sg.num_edges() == len(eids)
assert F.array_equal(sg.edata[dgl.EID], eids)
# Test init node data
new_shape = (g.number_of_nodes(), 2)
test1 = dgl.distributed.DistTensor(new_shape, F.int32)
......@@ -485,6 +499,14 @@ def check_dist_graph_hetero(g, num_clients, num_nodes, num_edges):
feats = F.squeeze(feats1, 1)
assert np.all(F.asnumpy(feats == eids))
# Test edge_subgraph
sg = g.edge_subgraph({'r1': eids})
assert sg.num_edges() == len(eids)
assert F.array_equal(sg.edata[dgl.EID], eids)
sg = g.edge_subgraph({('n1', 'r1', 'n2'): eids})
assert sg.num_edges() == len(eids)
assert F.array_equal(sg.edata[dgl.EID], eids)
# Test init node data
new_shape = (g.number_of_nodes('n1'), 2)
g.nodes['n1'].data['test1'] = dgl.distributed.DistTensor(new_shape, F.int32)
......
tests/distributed/test_distributed_sampling.py
View file @ e9b624fe
......@@ -84,7 +84,7 @@ def start_get_degrees_client(rank, tmpdir, disable_shared_mem, nids=None):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_get_degrees.json', rank)
dgl.distributed.initialize("rpc_ip_config.txt", 1)
dgl.distributed.initialize("rpc_ip_config.txt")
dist_graph = DistGraph("test_get_degrees", gpb=gpb)
try:
in_deg = dist_graph.in_degrees(nids)
......@@ -160,9 +160,9 @@ def check_rpc_find_edges_shuffle(tmpdir, num_server):
def create_random_hetero(dense=False, empty=False):
num_nodes = {'n1': 210, 'n2': 200, 'n3': 220} if dense else \
{'n1': 1010, 'n2': 1000, 'n3': 1020}
etypes = [('n1', 'r1', 'n2'),
('n1', 'r2', 'n3'),
('n2', 'r3', 'n3')]
etypes = [('n1', 'r12', 'n2'),
('n1', 'r13', 'n3'),
('n2', 'r23', 'n3')]
edges = {}
random.seed(42)
for etype in etypes:
......@@ -195,9 +195,18 @@ def check_rpc_hetero_find_edges_shuffle(tmpdir, num_server):
time.sleep(1)
pserver_list.append(p)
eids = F.tensor(np.random.randint(g.number_of_edges('r1'), size=100))
u, v = g.find_edges(orig_eid['r1'][eids], etype='r1')
du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids, etype='r1')
eids = F.tensor(np.random.randint(g.num_edges('r12'), size=100))
expect_except = False
try:
_, _ = g.find_edges(orig_eid['r12'][eids], etype=('n1', 'r12'))
except:
expect_except = True
assert expect_except
u, v = g.find_edges(orig_eid['r12'][eids], etype='r12')
u1, v1 = g.find_edges(orig_eid['r12'][eids], etype=('n1', 'r12', 'n2'))
assert F.array_equal(u, u1)
assert F.array_equal(v, v1)
du, dv = start_find_edges_client(0, tmpdir, num_server > 1, eids, etype='r12')
du = orig_nid['n1'][du]
dv = orig_nid['n2'][dv]
assert F.array_equal(u, du)
......@@ -335,7 +344,8 @@ def start_hetero_sample_client(rank, tmpdir, disable_shared_mem, nodes):
return block, gpb
def start_hetero_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
nodes={'n3': [0, 10, 99, 66, 124, 208]}):
nodes={'n3': [0, 10, 99, 66, 124, 208]},
etype_sorted=False):
gpb = None
if disable_shared_mem:
_, _, _, gpb, _, _, _ = load_partition(tmpdir / 'test_sampling.json', rank)
......@@ -358,7 +368,7 @@ def start_hetero_etype_sample_client(rank, tmpdir, disable_shared_mem, fanout=3,
if gpb is None:
gpb = dist_graph.get_partition_book()
try:
sampled_graph = sample_etype_neighbors(dist_graph, nodes, dgl.ETYPE, fanout)
sampled_graph = sample_etype_neighbors(dist_graph, nodes, dgl.ETYPE, fanout, etype_sorted=etype_sorted)
block = dgl.to_block(sampled_graph, nodes)
block.edata[dgl.EID] = sampled_graph.edata[dgl.EID]
except Exception as e:
......@@ -461,7 +471,7 @@ def check_rpc_hetero_sampling_empty_shuffle(tmpdir, num_server):
assert block.number_of_edges() == 0
assert len(block.etypes) == len(g.etypes)
def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server, etype_sorted=False):
generate_ip_config("rpc_ip_config.txt", num_server, num_server)
g = create_random_hetero(dense=True)
......@@ -474,21 +484,22 @@ def check_rpc_hetero_etype_sampling_shuffle(tmpdir, num_server):
pserver_list = []
ctx = mp.get_context('spawn')
for i in range(num_server):
p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling'))
p = ctx.Process(target=start_server, args=(i, tmpdir, num_server > 1, 'test_sampling', ['csc', 'coo']))
p.start()
time.sleep(1)
pserver_list.append(p)
fanout = 3
block, gpb = start_hetero_etype_sample_client(0, tmpdir, num_server > 1, fanout,
nodes={'n3': [0, 10, 99, 66, 124, 208]})
nodes={'n3': [0, 10, 99, 66, 124, 208]},
etype_sorted=etype_sorted)
print("Done sampling")
for p in pserver_list:
p.join()
src, dst = block.edges(etype=('n1', 'r2', 'n3'))
src, dst = block.edges(etype=('n1', 'r13', 'n3'))
assert len(src) == 18
src, dst = block.edges(etype=('n2', 'r3', 'n3'))
src, dst = block.edges(etype=('n2', 'r23', 'n3'))
assert len(src) == 18
orig_nid_map = {ntype: F.zeros((g.number_of_nodes(ntype),), dtype=F.int64) for ntype in g.ntypes}
......@@ -832,6 +843,7 @@ def test_rpc_sampling_shuffle(num_server):
check_rpc_hetero_sampling_shuffle(Path(tmpdirname), num_server)
check_rpc_hetero_sampling_empty_shuffle(Path(tmpdirname), num_server)
check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server)
check_rpc_hetero_etype_sampling_shuffle(Path(tmpdirname), num_server, etype_sorted=True)
check_rpc_hetero_etype_sampling_empty_shuffle(Path(tmpdirname), num_server)
check_rpc_bipartite_sampling_empty(Path(tmpdirname), num_server)
check_rpc_bipartite_sampling_shuffle(Path(tmpdirname), num_server)
......
tests/pytorch/test_dataloader.py
View file @ e9b624fe
......@@ -5,11 +5,9 @@ import dgl.ops as OPS
import backend as F
import unittest
import torch
import torch.distributed as dist
from functools import partial
from torch.utils.data import DataLoader
from collections import defaultdict
from collections.abc import Iterator, Mapping
from itertools import product
from test_utils import parametrize_idtype
import pytest
......@@ -40,7 +38,7 @@ def test_cluster_gcn(num_workers):
def test_shadow(num_workers):
g = dgl.data.CoraFullDataset()[0]
sampler = dgl.dataloading.ShaDowKHopSampler([5, 10, 15])
dataloader = dgl.dataloading.NodeDataLoader(
dataloader = dgl.dataloading.DataLoader(
g, torch.arange(g.num_nodes()), sampler,
batch_size=5, shuffle=True, drop_last=False, num_workers=num_workers)
for i, (input_nodes, output_nodes, subgraph) in enumerate(dataloader):
......@@ -70,43 +68,73 @@ def test_saint(num_workers, mode):
for sg in dataloader:
pass
@pytest.mark.parametrize('num_workers', [0, 4])
def test_neighbor_nonuniform(num_workers):
g = dgl.graph(([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]))
@parametrize_idtype
@pytest.mark.parametrize('mode', ['cpu', 'uva_cuda_indices', 'uva_cpu_indices', 'pure_gpu'])
@pytest.mark.parametrize('use_ddp', [False, True])
def test_neighbor_nonuniform(idtype, mode, use_ddp):
if mode != 'cpu' and F.ctx() == F.cpu():
pytest.skip('UVA and GPU sampling require a GPU.')
if use_ddp:
dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
'tcp://127.0.0.1:12347', world_size=1, rank=0)
g = dgl.graph(([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1])).astype(idtype)
g.edata['p'] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
if mode in ('cpu', 'uva_cpu_indices'):
indices = F.copy_to(F.tensor([0, 1], idtype), F.cpu())
else:
indices = F.copy_to(F.tensor([0, 1], idtype), F.cuda())
if mode == 'pure_gpu':
g = g.to(F.cuda())
use_uva = mode.startswith('uva')
sampler = dgl.dataloading.MultiLayerNeighborSampler([2], prob='p')
dataloader = dgl.dataloading.NodeDataLoader(g, [0, 1], sampler, batch_size=1, device=F.ctx())
for input_nodes, output_nodes, blocks in dataloader:
seed = output_nodes.item()
neighbors = set(input_nodes[1:].cpu().numpy())
if seed == 1:
assert neighbors == {5, 6}
elif seed == 0:
assert neighbors == {1, 2}
for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
dataloader = dgl.dataloading.NodeDataLoader(
g, indices, sampler,
batch_size=1, device=F.ctx(),
num_workers=num_workers,
use_uva=use_uva,
use_ddp=use_ddp)
for input_nodes, output_nodes, blocks in dataloader:
seed = output_nodes.item()
neighbors = set(input_nodes[1:].cpu().numpy())
if seed == 1:
assert neighbors == {5, 6}
elif seed == 0:
assert neighbors == {1, 2}
g = dgl.heterograph({
('B', 'BA', 'A'): ([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]),
('C', 'CA', 'A'): ([1, 2, 3, 4, 5, 6, 7, 8], [0, 0, 0, 0, 1, 1, 1, 1]),
})
}).astype(idtype)
g.edges['BA'].data['p'] = torch.FloatTensor([1, 1, 0, 0, 1, 1, 0, 0])
g.edges['CA'].data['p'] = torch.FloatTensor([0, 0, 1, 1, 0, 0, 1, 1])
sampler = dgl.dataloading.MultiLayerNeighborSampler([2], prob='p')
dataloader = dgl.dataloading.NodeDataLoader(
g, {'A': [0, 1]}, sampler, batch_size=1, device=F.ctx())
for input_nodes, output_nodes, blocks in dataloader:
seed = output_nodes['A'].item()
# Seed and neighbors are of different node types so slicing is not necessary here.
neighbors = set(input_nodes['B'].cpu().numpy())
if seed == 1:
assert neighbors == {5, 6}
elif seed == 0:
assert neighbors == {1, 2}
neighbors = set(input_nodes['C'].cpu().numpy())
if seed == 1:
assert neighbors == {7, 8}
elif seed == 0:
assert neighbors == {3, 4}
if mode == 'pure_gpu':
g = g.to(F.cuda())
for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
dataloader = dgl.dataloading.NodeDataLoader(
g, {'A': indices}, sampler,
batch_size=1, device=F.ctx(),
num_workers=num_workers,
use_uva=use_uva,
use_ddp=use_ddp)
for input_nodes, output_nodes, blocks in dataloader:
seed = output_nodes['A'].item()
# Seed and neighbors are of different node types so slicing is not necessary here.
neighbors = set(input_nodes['B'].cpu().numpy())
if seed == 1:
assert neighbors == {5, 6}
elif seed == 0:
assert neighbors == {1, 2}
neighbors = set(input_nodes['C'].cpu().numpy())
if seed == 1:
assert neighbors == {7, 8}
elif seed == 0:
assert neighbors == {3, 4}
if use_ddp:
dist.destroy_process_group()
def _check_dtype(data, dtype, attr_name):
if isinstance(data, dict):
......@@ -130,38 +158,37 @@ def _check_device(data):
@parametrize_idtype
@pytest.mark.parametrize('sampler_name', ['full', 'neighbor', 'neighbor2'])
@pytest.mark.parametrize('pin_graph', [None, 'cuda_indices', 'cpu_indices'])
def test_node_dataloader(idtype, sampler_name, pin_graph):
@pytest.mark.parametrize('mode', ['cpu', 'uva_cuda_indices', 'uva_cpu_indices', 'pure_gpu'])
@pytest.mark.parametrize('use_ddp', [False, True])
def test_node_dataloader(idtype, sampler_name, mode, use_ddp):
if mode != 'cpu' and F.ctx() == F.cpu():
pytest.skip('UVA and GPU sampling require a GPU.')
if use_ddp:
dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
'tcp://127.0.0.1:12347', world_size=1, rank=0)
g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
g1.ndata['feat'] = F.copy_to(F.randn((5, 8)), F.cpu())
g1.ndata['label'] = F.copy_to(F.randn((g1.num_nodes(),)), F.cpu())
indices = F.arange(0, g1.num_nodes(), idtype)
if F.ctx() != F.cpu():
if pin_graph:
g1.create_formats_()
g1.pin_memory_()
if pin_graph == 'cpu_indices':
indices = F.arange(0, g1.num_nodes(), idtype, F.cpu())
elif pin_graph == 'cuda_indices':
if F._default_context_str == 'gpu':
indices = F.arange(0, g1.num_nodes(), idtype, F.cuda())
else:
return # skip
else:
g1 = g1.to('cuda')
if mode in ('cpu', 'uva_cpu_indices'):
indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cpu())
else:
indices = F.copy_to(F.arange(0, g1.num_nodes(), idtype), F.cuda())
if mode == 'pure_gpu':
g1 = g1.to(F.cuda())
use_uva = pin_graph is not None and F.ctx() != F.cpu()
use_uva = mode.startswith('uva')
for num_workers in [0, 1, 2]:
sampler = {
'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
dataloader = dgl.dataloading.NodeDataLoader(
sampler = {
'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3]),
'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
dataloader = dgl.dataloading.DataLoader(
g1, indices, sampler, device=F.ctx(),
batch_size=g1.num_nodes(),
num_workers=(num_workers if (pin_graph and F.ctx() == F.cpu()) else 0),
use_uva=use_uva)
num_workers=num_workers,
use_uva=use_uva,
use_ddp=use_ddp)
for input_nodes, output_nodes, blocks in dataloader:
_check_device(input_nodes)
_check_device(output_nodes)
......@@ -169,8 +196,6 @@ def test_node_dataloader(idtype, sampler_name, pin_graph):
_check_dtype(input_nodes, idtype, 'dtype')
_check_dtype(output_nodes, idtype, 'dtype')
_check_dtype(blocks, idtype, 'idtype')
if g1.is_pinned():
g1.unpin_memory_()
g2 = dgl.heterograph({
('user', 'follow', 'user'): ([0, 0, 0, 1, 1, 1, 2], [1, 2, 3, 0, 2, 3, 0]),
......@@ -180,73 +205,80 @@ def test_node_dataloader(idtype, sampler_name, pin_graph):
}).astype(idtype)
for ntype in g2.ntypes:
g2.nodes[ntype].data['feat'] = F.copy_to(F.randn((g2.num_nodes(ntype), 8)), F.cpu())
indices = {nty: F.arange(0, g2.num_nodes(nty)) for nty in g2.ntypes}
if F.ctx() != F.cpu():
if pin_graph:
g2.create_formats_()
g2.pin_memory_()
if pin_graph == 'cpu_indices':
indices = {nty: F.arange(0, g2.num_nodes(nty), idtype, F.cpu()) for nty in g2.ntypes}
elif pin_graph == 'cuda_indices':
if F._default_context_str == 'gpu':
indices = {nty: F.arange(0, g2.num_nodes(), idtype, F.cuda()) for nty in g2.ntypes}
else:
return # skip
else:
g2 = g2.to('cuda')
if mode in ('cpu', 'uva_cpu_indices'):
indices = {nty: F.copy_to(g2.nodes(nty), F.cpu()) for nty in g2.ntypes}
else:
indices = {nty: F.copy_to(g2.nodes(nty), F.cuda()) for nty in g2.ntypes}
if mode == 'pure_gpu':
g2 = g2.to(F.cuda())
batch_size = max(g2.num_nodes(nty) for nty in g2.ntypes)
sampler = {
'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
'neighbor': dgl.dataloading.MultiLayerNeighborSampler([{etype: 3 for etype in g2.etypes}] * 2),
'neighbor2': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
for num_workers in [0, 1, 2] if mode == 'cpu' else [0]:
dataloader = dgl.dataloading.DataLoader(
g2, indices, sampler,
device=F.ctx(), batch_size=batch_size,
num_workers=num_workers,
use_uva=use_uva,
use_ddp=use_ddp)
assert isinstance(iter(dataloader), Iterator)
for input_nodes, output_nodes, blocks in dataloader:
_check_device(input_nodes)
_check_device(output_nodes)
_check_device(blocks)
_check_dtype(input_nodes, idtype, 'dtype')
_check_dtype(output_nodes, idtype, 'dtype')
_check_dtype(blocks, idtype, 'idtype')
dataloader = dgl.dataloading.NodeDataLoader(
g2, {nty: g2.nodes(nty) for nty in g2.ntypes},
sampler, device=F.ctx(), batch_size=batch_size,
num_workers=(num_workers if (pin_graph and F.ctx() == F.cpu()) else 0),
use_uva=use_uva)
assert isinstance(iter(dataloader), Iterator)
for input_nodes, output_nodes, blocks in dataloader:
_check_device(input_nodes)
_check_device(output_nodes)
_check_device(blocks)
_check_dtype(input_nodes, idtype, 'dtype')
_check_dtype(output_nodes, idtype, 'dtype')
_check_dtype(blocks, idtype, 'idtype')
if g2.is_pinned():
g2.unpin_memory_()
if use_ddp:
dist.destroy_process_group()
@parametrize_idtype
@pytest.mark.parametrize('sampler_name', ['full', 'neighbor'])
@pytest.mark.parametrize('neg_sampler', [
dgl.dataloading.negative_sampler.Uniform(2),
dgl.dataloading.negative_sampler.GlobalUniform(15, False, 3),
dgl.dataloading.negative_sampler.GlobalUniform(15, True, 3)])
@pytest.mark.parametrize('pin_graph', [False, True])
def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4]))
if F.ctx() != F.cpu() and pin_graph:
g1.create_formats_()
g1.pin_memory_()
@pytest.mark.parametrize('mode', ['cpu', 'uva', 'pure_gpu'])
@pytest.mark.parametrize('use_ddp', [False, True])
def test_edge_dataloader(idtype, sampler_name, neg_sampler, mode, use_ddp):
if mode != 'cpu' and F.ctx() == F.cpu():
pytest.skip('UVA and GPU sampling require a GPU.')
if mode == 'uva' and isinstance(neg_sampler, dgl.dataloading.negative_sampler.GlobalUniform):
pytest.skip("GlobalUniform don't support UVA yet.")
if use_ddp:
dist.init_process_group('gloo' if F.ctx() == F.cpu() else 'nccl',
'tcp://127.0.0.1:12347', world_size=1, rank=0)
g1 = dgl.graph(([0, 0, 0, 1, 1], [1, 2, 3, 3, 4])).astype(idtype)
g1.ndata['feat'] = F.copy_to(F.randn((5, 8)), F.cpu())
if mode == 'pure_gpu':
g1 = g1.to(F.cuda())
sampler = {
'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
'neighbor': dgl.dataloading.MultiLayerNeighborSampler([3, 3])}[sampler_name]
# no negative sampler
dataloader = dgl.dataloading.EdgeDataLoader(
g1, g1.edges(form='eid'), sampler, device=F.ctx(), batch_size=g1.num_edges())
edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
dataloader = dgl.dataloading.DataLoader(
g1, g1.edges(form='eid'), edge_sampler,
device=F.ctx(), batch_size=g1.num_edges(),
use_uva=(mode == 'uva'), use_ddp=use_ddp)
for input_nodes, pos_pair_graph, blocks in dataloader:
_check_device(input_nodes)
_check_device(pos_pair_graph)
_check_device(blocks)
# negative sampler
dataloader = dgl.dataloading.EdgeDataLoader(
g1, g1.edges(form='eid'), sampler, device=F.ctx(),
negative_sampler=neg_sampler, batch_size=g1.num_edges())
edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
sampler, negative_sampler=neg_sampler)
dataloader = dgl.dataloading.DataLoader(
g1, g1.edges(form='eid'), edge_sampler,
device=F.ctx(), batch_size=g1.num_edges(),
use_uva=(mode == 'uva'), use_ddp=use_ddp)
for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
_check_device(input_nodes)
_check_device(pos_pair_graph)
......@@ -258,9 +290,12 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
('user', 'followed-by', 'user'): ([1, 2, 3, 0, 2, 3, 0], [0, 0, 0, 1, 1, 1, 2]),
('user', 'play', 'game'): ([0, 1, 1, 3, 5], [0, 1, 2, 0, 2]),
('game', 'played-by', 'user'): ([0, 1, 2, 0, 2], [0, 1, 1, 3, 5])
})
}).astype(idtype)
for ntype in g2.ntypes:
g2.nodes[ntype].data['feat'] = F.copy_to(F.randn((g2.num_nodes(ntype), 8)), F.cpu())
if mode == 'pure_gpu':
g2 = g2.to(F.cuda())
batch_size = max(g2.num_edges(ety) for ety in g2.canonical_etypes)
sampler = {
'full': dgl.dataloading.MultiLayerFullNeighborSampler(2),
......@@ -268,19 +303,23 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
}[sampler_name]
# no negative sampler
dataloader = dgl.dataloading.EdgeDataLoader(
edge_sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
dataloader = dgl.dataloading.DataLoader(
g2, {ety: g2.edges(form='eid', etype=ety) for ety in g2.canonical_etypes},
sampler, device=F.ctx(), batch_size=batch_size)
edge_sampler, device=F.ctx(), batch_size=batch_size,
use_uva=(mode == 'uva'), use_ddp=use_ddp)
for input_nodes, pos_pair_graph, blocks in dataloader:
_check_device(input_nodes)
_check_device(pos_pair_graph)
_check_device(blocks)
# negative sampler
dataloader = dgl.dataloading.EdgeDataLoader(
edge_sampler = dgl.dataloading.as_edge_prediction_sampler(
sampler, negative_sampler=neg_sampler)
dataloader = dgl.dataloading.DataLoader(
g2, {ety: g2.edges(form='eid', etype=ety) for ety in g2.canonical_etypes},
sampler, device=F.ctx(), negative_sampler=neg_sampler,
batch_size=batch_size)
edge_sampler, device=F.ctx(),batch_size=batch_size,
use_uva=(mode == 'uva'), use_ddp=use_ddp)
assert isinstance(iter(dataloader), Iterator)
for input_nodes, pos_pair_graph, neg_pair_graph, blocks in dataloader:
......@@ -289,8 +328,8 @@ def test_edge_dataloader(sampler_name, neg_sampler, pin_graph):
_check_device(neg_pair_graph)
_check_device(blocks)
if g1.is_pinned():
g1.unpin_memory_()
if use_ddp:
dist.destroy_process_group()
def _create_homogeneous():
s = torch.randint(0, 200, (1000,), device=F.ctx())
......
tests/pytorch/test_nn.py
View file @ e9b624fe
......@@ -412,6 +412,60 @@ def test_rgcn(idtype, O):
h_new = rgc_bdd(g, h, r, norm)
assert h_new.shape == (100, O)
@parametrize_idtype
@pytest.mark.parametrize('O', [1, 10, 40])
def test_rgcn_default_nbasis(idtype, O):
ctx = F.ctx()
etype = []
g = dgl.from_scipy(sp.sparse.random(100, 100, density=0.1))
g = g.astype(idtype).to(F.ctx())
# 5 etypes
R = 5
for i in range(g.number_of_edges()):
etype.append(i % 5)
I = 10
h = th.randn((100, I)).to(ctx)
r = th.tensor(etype).to(ctx)
norm = th.rand((g.number_of_edges(), 1)).to(ctx)
sorted_r, idx = th.sort(r)
sorted_g = dgl.reorder_graph(g, edge_permute_algo='custom', permute_config={'edges_perm' : idx.to(idtype)})
sorted_norm = norm[idx]
rgc = nn.RelGraphConv(I, O, R).to(ctx)
th.save(rgc, tmp_buffer) # test pickle
rgc_basis = nn.RelGraphConv(I, O, R, "basis").to(ctx)
th.save(rgc_basis, tmp_buffer) # test pickle
if O % R == 0:
rgc_bdd = nn.RelGraphConv(I, O, R, "bdd").to(ctx)
th.save(rgc_bdd, tmp_buffer) # test pickle
# basic usage
h_new = rgc(g, h, r)
assert h_new.shape == (100, O)
h_new_basis = rgc_basis(g, h, r)
assert h_new_basis.shape == (100, O)
if O % R == 0:
h_new_bdd = rgc_bdd(g, h, r)
assert h_new_bdd.shape == (100, O)
# sorted input
h_new_sorted = rgc(sorted_g, h, sorted_r, presorted=True)
assert th.allclose(h_new, h_new_sorted, atol=1e-4, rtol=1e-4)
h_new_basis_sorted = rgc_basis(sorted_g, h, sorted_r, presorted=True)
assert th.allclose(h_new_basis, h_new_basis_sorted, atol=1e-4, rtol=1e-4)
if O % R == 0:
h_new_bdd_sorted = rgc_bdd(sorted_g, h, sorted_r, presorted=True)
assert th.allclose(h_new_bdd, h_new_bdd_sorted, atol=1e-4, rtol=1e-4)
# norm input
h_new = rgc(g, h, r, norm)
assert h_new.shape == (100, O)
h_new = rgc_basis(g, h, r, norm)
assert h_new.shape == (100, O)
if O % R == 0:
h_new = rgc_bdd(g, h, r, norm)
assert h_new.shape == (100, O)
@parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
......@@ -647,10 +701,11 @@ def test_appnp_conv_e_weight(g, idtype):
@parametrize_idtype
@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
def test_gcn2conv_e_weight(g, idtype):
@pytest.mark.parametrize("bias", [True, False])
def test_gcn2conv_e_weight(g, idtype, bias):
ctx = F.ctx()
g = g.astype(idtype).to(ctx)
gcn2conv = nn.GCN2Conv(5, layer=2, alpha=0.5,
gcn2conv = nn.GCN2Conv(5, layer=2, alpha=0.5, bias=bias,
project_initial_features=True)
feat = F.randn((g.number_of_nodes(), 5))
eweight = F.ones((g.num_edges(), ))
......@@ -1207,6 +1262,11 @@ def test_hetero_embedding(out_dim):
assert embeds['user'].shape == (2, out_dim)
assert embeds[('user', 'follows', 'user')].shape == (3, out_dim)
layer.reset_parameters()
embeds = layer.weight
assert embeds['user'].shape == (2, out_dim)
assert embeds[('user', 'follows', 'user')].shape == (3, out_dim)
embeds = layer({
'user': F.tensor([0], dtype=F.int64),
('user', 'follows', 'user'): F.tensor([0, 2], dtype=F.int64)
......@@ -1561,4 +1621,4 @@ def test_dgn_conv(in_size, out_size, aggregators, scalers, delta,
aggregators_non_eig = [aggr for aggr in aggregators if not aggr.startswith('dir')]
model = nn.DGNConv(in_size, out_size, aggregators_non_eig, scalers, delta, dropout,
num_towers, edge_feat_size, residual).to(dev)
model(g, h, edge_feat=e)
model(g, h, edge_feat=e)
\ No newline at end of file
tests/pytorch/test_unified_tensor.py
View file @ e9b624fe
......@@ -18,23 +18,24 @@ def start_unified_tensor_worker(dev_id, input, seq_idx, rand_idx, output_seq, ou
def test_unified_tensor():
test_row_size = 65536
test_col_size = 128
rand_test_size = 8192
device = th.device('cuda:0')
input = th.rand((test_row_size, test_col_size))
input_unified = dgl.contrib.UnifiedTensor(input, device=th.device('cuda'))
input_unified = dgl.contrib.UnifiedTensor(input, device=device)
seq_idx = th.arange(0, test_row_size)
# CPU indexing
assert th.all(th.eq(input[seq_idx], input_unified[seq_idx]))
seq_idx = seq_idx.to(th.device('cuda'))
assert th.all(th.eq(input[seq_idx].to(th.device('cuda')), input_unified[seq_idx]))
# GPU indexing
assert th.all(th.eq(input[seq_idx].to(device), input_unified[seq_idx.to(device)]))
rand_idx = th.randint(0, test_row_size, (rand_test_size,))
# CPU indexing
assert th.all(th.eq(input[rand_idx], input_unified[rand_idx]))
rand_idx = rand_idx.to(th.device('cuda'))
assert th.all(th.eq(input[rand_idx].to(th.device('cuda')), input_unified[rand_idx]))
# GPU indexing
assert th.all(th.eq(input[rand_idx].to(device), input_unified[rand_idx.to(device)]))
@unittest.skipIf(os.name == 'nt', reason='Do not support windows yet')
@unittest.skipIf(F.ctx().type == 'cpu', reason='gpu only test')
......
tests/scripts/build_dgl.sh
View file @ e9b624fe
......@@ -3,7 +3,7 @@ set -e
. /opt/conda/etc/profile.d/conda.sh
if [ $# -ne 1 ]; then
echo "Device argument required, can be cpu or gpu"
echo "Device argument required, can be cpu, gpu or cugraph"
exit -1
fi
......@@ -11,7 +11,9 @@ CMAKE_VARS="-DBUILD_CPP_TEST=ON -DUSE_OPENMP=ON -DBUILD_TORCH=ON"
# This is a semicolon-separated list of Python interpreters containing PyTorch.
# The value here is for CI. Replace it with your own or comment this whole
# statement for default Python interpreter.
CMAKE_VARS="$CMAKE_VARS -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
if [ "$1" != "cugraph" ]; then
CMAKE_VARS="$CMAKE_VARS -DTORCH_PYTHON_INTERPS=/opt/conda/envs/pytorch-ci/bin/python"
fi
#This is implemented to detect underlying architecture and enable arch specific optimization.
arch=`uname -m`
......@@ -19,7 +21,7 @@ if [[ $arch == *"x86"* ]]; then
CMAKE_VARS="-DUSE_AVX=ON $CMAKE_VARS"
fi
if [ "$1" == "gpu" ]; then
if [[ $1 != "cpu" ]]; then
CMAKE_VARS="-DUSE_CUDA=ON -DUSE_NCCL=ON -DUSE_FP16=ON $CMAKE_VARS"
fi
......@@ -36,14 +38,23 @@ make -j
popd
pushd python
for backend in pytorch mxnet tensorflow
do
conda activate "${backend}-ci"
rm -rf build *.egg-info dist
pip uninstall -y dgl
# test install
python3 setup.py install
# test inplace build (for cython)
python3 setup.py build_ext --inplace
done
if [[ $1 == "cugraph" ]]; then
rm -rf build *.egg-info dist
pip uninstall -y dgl
# test install
python3 setup.py install
# test inplace build (for cython)
python3 setup.py build_ext --inplace
else
for backend in pytorch mxnet tensorflow
do
conda activate "${backend}-ci"
rm -rf build *.egg-info dist
pip uninstall -y dgl
# test install
python3 setup.py install
# test inplace build (for cython)
python3 setup.py build_ext --inplace
done
fi
popd
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