"host/online_compile/include/logger.hpp" did not exist on "d2315b0dfcd6f31cca4328819eaf60d77e952dd6"
test_fps.py 2.31 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
from torch import Tensor
rusty1s's avatar
rusty1s committed
6
from torch_cluster import fps
Matthias Fey's avatar
Matthias Fey committed
7
from torch_cluster.testing import devices, grad_dtypes, tensor
rusty1s's avatar
rusty1s committed
8
9


rusty1s's avatar
rusty1s committed
10
11
12
13
14
@torch.jit.script
def fps2(x: Tensor, ratio: Tensor) -> Tensor:
    return fps(x, None, ratio, False)


rusty1s's avatar
rusty1s committed
15
@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
rusty1s's avatar
rusty1s committed
16
def test_fps(dtype, device):
rusty1s's avatar
rusty1s committed
17
18
19
20
21
22
23
24
25
26
    x = tensor([
        [-1, -1],
        [-1, +1],
        [+1, +1],
        [+1, -1],
        [-2, -2],
        [-2, +2],
        [+2, +2],
        [+2, -2],
    ], dtype, device)
rusty1s's avatar
rusty1s committed
27
    batch = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
28
29
    ptr_list = [0, 4, 8]
    ptr = torch.tensor(ptr_list, device=device)
rusty1s's avatar
rusty1s committed
30

rusty1s's avatar
rusty1s committed
31
32
33
    out = fps(x, batch, random_start=False)
    assert out.tolist() == [0, 2, 4, 6]

Duc's avatar
Duc committed
34
35
36
    out = fps(x, batch, ratio=0.5, random_start=False)
    assert out.tolist() == [0, 2, 4, 6]

37
38
    ratio = torch.tensor(0.5, device=device)
    out = fps(x, batch, ratio=ratio, random_start=False)
rusty1s's avatar
rusty1s committed
39
    assert out.tolist() == [0, 2, 4, 6]
rusty1s's avatar
rusty1s committed
40

41
42
43
44
45
46
47
48
    out = fps(x, ptr=ptr_list, ratio=0.5, random_start=False)
    assert out.tolist() == [0, 2, 4, 6]

    out = fps(x, ptr=ptr, ratio=0.5, random_start=False)
    assert out.tolist() == [0, 2, 4, 6]

    ratio = torch.tensor([0.5, 0.5], device=device)
    out = fps(x, batch, ratio=ratio, random_start=False)
Duc's avatar
Duc committed
49
50
    assert out.tolist() == [0, 2, 4, 6]

rusty1s's avatar
rusty1s committed
51
52
53
54
55
56
57
58
59
60
    out = fps(x, random_start=False)
    assert out.sort()[0].tolist() == [0, 5, 6, 7]

    out = fps(x, ratio=0.5, random_start=False)
    assert out.sort()[0].tolist() == [0, 5, 6, 7]

    out = fps(x, ratio=torch.tensor(0.5, device=device), random_start=False)
    assert out.sort()[0].tolist() == [0, 5, 6, 7]

    out = fps(x, ratio=torch.tensor([0.5], device=device), random_start=False)
rusty1s's avatar
rusty1s committed
61
    assert out.sort()[0].tolist() == [0, 5, 6, 7]
rusty1s's avatar
rusty1s committed
62

rusty1s's avatar
rusty1s committed
63
64
65
    out = fps2(x, torch.tensor([0.5], device=device))
    assert out.sort()[0].tolist() == [0, 5, 6, 7]

rusty1s's avatar
rusty1s committed
66
67
68
69
70
71
72
73
74

@pytest.mark.parametrize('device', devices)
def test_random_fps(device):
    N = 1024
    for _ in range(5):
        pos = torch.randn((2 * N, 3), device=device)
        batch_1 = torch.zeros(N, dtype=torch.long, device=device)
        batch_2 = torch.ones(N, dtype=torch.long, device=device)
        batch = torch.cat([batch_1, batch_2])
rusty1s's avatar
rusty1s committed
75
        idx = fps(pos, batch, ratio=0.5)
rusty1s's avatar
rusty1s committed
76
        assert idx.min() >= 0 and idx.max() < 2 * N