test_weighting.py 1.73 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
from itertools import product

import pytest
import torch
rusty1s's avatar
rusty1s committed
5
6
7
from torch.autograd import gradcheck
from torch_spline_conv.weighting import SplineWeighting
from torch_spline_conv.basis import SplineBasis
rusty1s's avatar
rusty1s committed
8

rusty1s's avatar
rusty1s committed
9
from .utils import dtypes, devices, tensor
rusty1s's avatar
rusty1s committed
10

rusty1s's avatar
rusty1s committed
11
12
devices = [torch.device('cpu')]

rusty1s's avatar
rusty1s committed
13
tests = [{
rusty1s's avatar
rusty1s committed
14
    'x': [[1, 2], [3, 4]],
rusty1s's avatar
rusty1s committed
15
16
17
    'weight': [[[1], [2]], [[3], [4]], [[5], [6]], [[7], [8]]],
    'basis': [[0.5, 0, 0.5, 0], [0, 0, 0.5, 0.5]],
    'weight_index': [[0, 1, 2, 3], [0, 1, 2, 3]],
rusty1s's avatar
rename  
rusty1s committed
18
    'expected': [
rusty1s's avatar
rusty1s committed
19
20
21
22
23
24
        [0.5 * ((1 * (1 + 5)) + (2 * (2 + 6)))],
        [0.5 * ((3 * (5 + 7)) + (4 * (6 + 8)))],
    ]
}]


rusty1s's avatar
rusty1s committed
25
26
@pytest.mark.parametrize('test,dtype,device', product(tests, dtypes, devices))
def test_spline_weighting_forward(test, dtype, device):
rusty1s's avatar
rusty1s committed
27
    x = tensor(test['x'], dtype, device)
rusty1s's avatar
rusty1s committed
28
29
30
    weight = tensor(test['weight'], dtype, device)
    basis = tensor(test['basis'], dtype, device)
    weight_index = tensor(test['weight_index'], torch.long, device)
rusty1s's avatar
rusty1s committed
31

rusty1s's avatar
rusty1s committed
32
    out = SplineWeighting.apply(x, weight, basis, weight_index)
rusty1s's avatar
rename  
rusty1s committed
33
    assert out.tolist() == test['expected']
rusty1s's avatar
rusty1s committed
34
35


rusty1s's avatar
rusty1s committed
36
37
38
39
40
@pytest.mark.parametrize('device', devices)
def test_spline_basis_backward(device):
    pseudo = torch.rand((4, 2), dtype=torch.double, device=device)
    kernel_size = tensor([5, 5], torch.long, device)
    is_open_spline = tensor([1, 1], torch.uint8, device)
rusty1s's avatar
rusty1s committed
41
    degree = 1
rusty1s's avatar
rusty1s committed
42

rusty1s's avatar
rusty1s committed
43
44
45
    op = SplineBasis.apply
    basis, weight_index = op(pseudo, kernel_size, is_open_spline, degree)
    basis.requires_grad_()
rusty1s's avatar
rusty1s committed
46

rusty1s's avatar
rusty1s committed
47
48
    x = torch.rand((4, 2), dtype=torch.double, device=device)
    x.requires_grad_()
rusty1s's avatar
rusty1s committed
49
50
    weight = torch.rand((25, 2, 4), dtype=torch.double, device=device)
    weight.requires_grad_()
rusty1s's avatar
rusty1s committed
51

rusty1s's avatar
rusty1s committed
52
    data = (x, weight, basis, weight_index)
rusty1s's avatar
rusty1s committed
53
    assert gradcheck(SplineWeighting.apply, data, eps=1e-6, atol=1e-4) is True