dmlc--dgl
Tingyu Wang
ba21295cd5
* update agg function with new bindings * handle optional import in __init__ * raise error in RelGraphConvAgg when pylibcugraphops not imported * Update tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py Co-authored-by: Mufei Li <mufeili1996@gmail.com> * Update tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py Co-authored-by: Mufei Li <mufeili1996@gmail.com> * use keyword args for readability * add missing docstring to pass CI * catch ImportError rather than ModuleNotFoundError Co-authored-by: Mufei Li <mufeili1996@gmail.com>
99 行
3.1 KiB
Python
99 行
3.1 KiB
Python
import pytest
|
|
import torch
|
|
import dgl
|
|
from dgl.nn import CuGraphRelGraphConv
|
|
from dgl.nn import RelGraphConv
|
|
|
|
# TODO(tingyu66): Re-enable the following tests after updating cuGraph CI image.
|
|
use_longs = [False, True]
|
|
max_in_degrees = [None, 8]
|
|
regularizers = [None, "basis"]
|
|
device = "cuda"
|
|
|
|
|
|
def generate_graph():
|
|
u = torch.tensor([0, 1, 0, 2, 3, 0, 4, 0, 5, 0, 6, 7, 0, 8, 9])
|
|
v = torch.tensor([1, 9, 2, 9, 9, 4, 9, 5, 9, 6, 9, 9, 8, 9, 0])
|
|
g = dgl.graph((u, v))
|
|
num_rels = 3
|
|
g.edata[dgl.ETYPE] = torch.randint(num_rels, (g.num_edges(),))
|
|
return g
|
|
|
|
@pytest.mark.skip()
|
|
@pytest.mark.parametrize('use_long', use_longs)
|
|
@pytest.mark.parametrize('max_in_degree', max_in_degrees)
|
|
@pytest.mark.parametrize("regularizer", regularizers)
|
|
def test_full_graph(use_long, max_in_degree, regularizer):
|
|
in_feat, out_feat, num_rels, num_bases = 10, 2, 3, 2
|
|
kwargs = {
|
|
"num_bases": num_bases,
|
|
"regularizer": regularizer,
|
|
"bias": False,
|
|
"self_loop": False,
|
|
}
|
|
g = generate_graph().to(device)
|
|
if use_long:
|
|
g = g.long()
|
|
else:
|
|
g = g.int()
|
|
feat = torch.ones(g.num_nodes(), in_feat).to(device)
|
|
|
|
torch.manual_seed(0)
|
|
conv1 = RelGraphConv(in_feat, out_feat, num_rels, **kwargs).to(device)
|
|
|
|
torch.manual_seed(0)
|
|
conv2 = CuGraphRelGraphConv(
|
|
in_feat, out_feat, num_rels, max_in_degree=max_in_degree, **kwargs
|
|
).to(device)
|
|
|
|
out1 = conv1(g, feat, g.edata[dgl.ETYPE])
|
|
out2 = conv2(g, feat, g.edata[dgl.ETYPE])
|
|
|
|
assert torch.allclose(out1, out2, atol=1e-06)
|
|
|
|
grad_out = torch.rand_like(out1)
|
|
out1.backward(grad_out)
|
|
out2.backward(grad_out)
|
|
assert torch.allclose(conv1.linear_r.W.grad, conv2.W.grad, atol=1e-6)
|
|
if regularizer is not None:
|
|
assert torch.allclose(
|
|
conv1.linear_r.coeff.grad, conv2.coeff.grad, atol=1e-6
|
|
)
|
|
|
|
@pytest.mark.skip()
|
|
@pytest.mark.parametrize('max_in_degree', max_in_degrees)
|
|
@pytest.mark.parametrize("regularizer", regularizers)
|
|
def test_mfg(max_in_degree, regularizer):
|
|
in_feat, out_feat, num_rels, num_bases = 10, 2, 3, 2
|
|
kwargs = {
|
|
"num_bases": num_bases,
|
|
"regularizer": regularizer,
|
|
"bias": False,
|
|
"self_loop": False,
|
|
}
|
|
g = generate_graph().to(device)
|
|
block = dgl.to_block(g)
|
|
feat = torch.ones(g.num_nodes(), in_feat).to(device)
|
|
|
|
torch.manual_seed(0)
|
|
conv1 = RelGraphConv(in_feat, out_feat, num_rels, **kwargs).to(device)
|
|
|
|
torch.manual_seed(0)
|
|
conv2 = CuGraphRelGraphConv(
|
|
in_feat, out_feat, num_rels, max_in_degree=max_in_degree, **kwargs
|
|
).to(device)
|
|
|
|
out1 = conv1(block, feat[block.srcdata[dgl.NID]], block.edata[dgl.ETYPE])
|
|
out2 = conv2(block, feat[block.srcdata[dgl.NID]], block.edata[dgl.ETYPE])
|
|
|
|
assert torch.allclose(out1, out2, atol=1e-06)
|
|
|
|
grad_out = torch.rand_like(out1)
|
|
out1.backward(grad_out)
|
|
out2.backward(grad_out)
|
|
assert torch.allclose(conv1.linear_r.W.grad, conv2.W.grad, atol=1e-6)
|
|
if regularizer is not None:
|
|
assert torch.allclose(
|
|
conv1.linear_r.coeff.grad, conv2.coeff.grad, atol=1e-6
|
|
)
|