test_roi_align.py 3.88 KB
Newer Older
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import numpy as np
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
3
import pytest
4
5
import torch

6
7
8
from mmcv.device.mlu import IS_MLU_AVAILABLE
from mmcv.utils import IS_CUDA_AVAILABLE

9
10
11
12
13
14
15
_USING_PARROTS = True
try:
    from parrots.autograd import gradcheck
except ImportError:
    from torch.autograd import gradcheck
    _USING_PARROTS = False

Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
16
# yapf:disable
17

Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
18
19
20
21
22
23
24
25
26
27
inputs = [([[[[1., 2.], [3., 4.]]]],
           [[0., 0., 0., 1., 1.]]),
          ([[[[1., 2.], [3., 4.]],
             [[4., 3.], [2., 1.]]]],
           [[0., 0., 0., 1., 1.]]),
          ([[[[1., 2., 5., 6.], [3., 4., 7., 8.],
              [9., 10., 13., 14.], [11., 12., 15., 16.]]]],
           [[0., 0., 0., 3., 3.]])]
outputs = [([[[[1.0, 1.25], [1.5, 1.75]]]],
            [[[[3.0625, 0.4375], [0.4375, 0.0625]]]]),
28
           ([[[[1.0, 1.25], [1.5, 1.75]],
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
29
30
31
              [[4.0, 3.75], [3.5, 3.25]]]],
            [[[[3.0625, 0.4375], [0.4375, 0.0625]],
              [[3.0625, 0.4375], [0.4375, 0.0625]]]]),
32
33
34
35
36
           ([[[[1.9375, 4.75], [7.5625, 10.375]]]],
            [[[[0.47265625, 0.42968750, 0.42968750, 0.04296875],
               [0.42968750, 0.39062500, 0.39062500, 0.03906250],
               [0.42968750, 0.39062500, 0.39062500, 0.03906250],
               [0.04296875, 0.03906250, 0.03906250, 0.00390625]]]])]
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
37
# yapf:enable
38

Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
39
40
41
42
pool_h = 2
pool_w = 2
spatial_scale = 1.0
sampling_ratio = 2
43
44


Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
45
46
def _test_roialign_gradcheck(device, dtype):
    try:
47
        from mmcv.ops import RoIAlign
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
48
49
50
51
52
53
54
    except ModuleNotFoundError:
        pytest.skip('RoIAlign op is not successfully compiled')
    if dtype is torch.half:
        pytest.skip('grad check does not support fp16')
    for case in inputs:
        np_input = np.array(case[0])
        np_rois = np.array(case[1])
55

Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
56
57
58
        x = torch.tensor(
            np_input, dtype=dtype, device=device, requires_grad=True)
        rois = torch.tensor(np_rois, dtype=dtype, device=device)
59

Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
60
        froipool = RoIAlign((pool_h, pool_w), spatial_scale, sampling_ratio)
61

62
63
64
65
66
        if torch.__version__ == 'parrots':
            gradcheck(
                froipool, (x, rois), no_grads=[rois], delta=1e-5, pt_atol=1e-5)
        else:
            gradcheck(froipool, (x, rois), eps=1e-5, atol=1e-5)
67
68


Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
69
70
def _test_roialign_allclose(device, dtype):
    try:
71
        from mmcv.ops import roi_align
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
    except ModuleNotFoundError:
        pytest.skip('test requires compilation')
    pool_h = 2
    pool_w = 2
    spatial_scale = 1.0
    sampling_ratio = 2
    for case, output in zip(inputs, outputs):
        np_input = np.array(case[0])
        np_rois = np.array(case[1])
        np_output = np.array(output[0])
        np_grad = np.array(output[1])

        x = torch.tensor(
            np_input, dtype=dtype, device=device, requires_grad=True)
        rois = torch.tensor(np_rois, dtype=dtype, device=device)

        output = roi_align(x, rois, (pool_h, pool_w), spatial_scale,
                           sampling_ratio, 'avg', True)
        output.backward(torch.ones_like(output))
        assert np.allclose(
            output.data.type(torch.float).cpu().numpy(), np_output, atol=1e-3)
        assert np.allclose(
            x.grad.data.type(torch.float).cpu().numpy(), np_grad, atol=1e-3)


97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
@pytest.mark.parametrize('device', [
    'cpu',
    pytest.param(
        'cuda',
        marks=pytest.mark.skipif(
            not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
    pytest.param(
        'mlu',
        marks=pytest.mark.skipif(
            not IS_MLU_AVAILABLE, reason='requires MLU support'))
])
@pytest.mark.parametrize('dtype', [
    torch.float,
    pytest.param(
        torch.double,
        marks=pytest.mark.skipif(
            IS_MLU_AVAILABLE,
            reason='MLU does not support for 64-bit floating point')),
    torch.half
])
Jerry Jiarui XU's avatar
Jerry Jiarui XU committed
117
118
119
120
121
def test_roialign(device, dtype):
    # check double only
    if dtype is torch.double:
        _test_roialign_gradcheck(device=device, dtype=dtype)
    _test_roialign_allclose(device=device, dtype=dtype)