test_deform_conv.py 8.19 KB
Newer Older
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import numpy as np
3
import pytest
4
5
import torch

6
from mmcv.utils import TORCH_VERSION, digit_version
7
8
9
10
11
12
13
14

try:
    # If PyTorch version >= 1.6.0 and fp16 is enabled, torch.cuda.amp.autocast
    # would be imported and used; we should test if our modules support it.
    from torch.cuda.amp import autocast
except ImportError:
    pass

15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
input = [[[[1., 2., 3.], [0., 1., 2.], [3., 5., 2.]]]]
offset_weight = [[[0.1, 0.4, 0.6, 0.1]], [[0.3, 0.2, 0.1, 0.3]],
                 [[0.5, 0.5, 0.2, 0.8]], [[0.8, 0.3, 0.9, 0.1]],
                 [[0.3, 0.1, 0.2, 0.5]], [[0.3, 0.7, 0.5, 0.3]],
                 [[0.6, 0.2, 0.5, 0.3]], [[0.4, 0.1, 0.8, 0.4]]]
offset_bias = [0.7, 0.1, 0.8, 0.5, 0.6, 0.5, 0.4, 0.7]
deform_weight = [[[0.4, 0.2, 0.1, 0.9]]]

gt_out = [[[[1.650, 0.], [0.000, 0.]]]]
gt_x_grad = [[[[-0.666, 0.204, 0.000], [0.030, -0.416, 0.012],
               [0.000, 0.252, 0.129]]]]
gt_offset_weight_grad = [[[[1.44, 2.88], [0.00, 1.44]]],
                         [[[-0.72, -1.44], [0.00, -0.72]]],
                         [[[0.00, 0.00], [0.00, 0.00]]],
                         [[[0.00, 0.00], [0.00, 0.00]]],
                         [[[-0.10, -0.20], [0.00, -0.10]]],
                         [[[-0.08, -0.16], [0.00, -0.08]]],
                         [[[-0.54, -1.08], [0.00, -0.54]]],
                         [[[-0.54, -1.08], [0.00, -0.54]]]]
gt_offset_bias_grad = [1.44, -0.72, 0., 0., -0.10, -0.08, -0.54, -0.54],
gt_deform_weight_grad = [[[[3.62, 0.], [0.40, 0.18]]]]


class TestDeformconv(object):

40
41
42
    def _test_deformconv(self,
                         dtype=torch.float,
                         threshold=1e-3,
43
44
45
                         device='cuda',
                         batch_size=10,
                         im2col_step=2):
46
47
        if not torch.cuda.is_available() and device == 'cuda':
            pytest.skip('test requires GPU')
48
49
50
        from mmcv.ops import DeformConv2dPack
        c_in = 1
        c_out = 1
51
52
53
54
55
        batch_size = 10
        repeated_input = np.repeat(input, batch_size, axis=0)
        repeated_gt_out = np.repeat(gt_out, batch_size, axis=0)
        repeated_gt_x_grad = np.repeat(gt_x_grad, batch_size, axis=0)
        x = torch.tensor(repeated_input, device=device, dtype=dtype)
56
        x.requires_grad = True
57
58
59
60
61
62
63
        model = DeformConv2dPack(
            in_channels=c_in,
            out_channels=c_out,
            kernel_size=2,
            stride=1,
            padding=0,
            im2col_step=im2col_step)
64
65
66
67
68
69
        model.conv_offset.weight.data = torch.nn.Parameter(
            torch.Tensor(offset_weight).reshape(8, 1, 2, 2))
        model.conv_offset.bias.data = torch.nn.Parameter(
            torch.Tensor(offset_bias).reshape(8))
        model.weight.data = torch.nn.Parameter(
            torch.Tensor(deform_weight).reshape(1, 1, 2, 2))
70
71
72
        if device == 'cuda':
            model.cuda()
        model.type(dtype)
73
74
75
76

        out = model(x)
        out.backward(torch.ones_like(out))

77
78
79
80
81
82
        assert np.allclose(out.data.detach().cpu().numpy(), repeated_gt_out,
                           threshold)
        assert np.allclose(x.grad.detach().cpu().numpy(), repeated_gt_x_grad,
                           threshold)
        # the batch size of the input is increased which results in
        # a larger gradient so we need to divide by the batch_size
83
        assert np.allclose(
84
            model.conv_offset.weight.grad.detach().cpu().numpy() / batch_size,
85
            gt_offset_weight_grad, threshold)
86
87
88
89
90
91
        assert np.allclose(
            model.conv_offset.bias.grad.detach().cpu().numpy() / batch_size,
            gt_offset_bias_grad, threshold)
        assert np.allclose(
            model.weight.grad.detach().cpu().numpy() / batch_size,
            gt_deform_weight_grad, threshold)
92

93
        from mmcv.ops import DeformConv2d
94

95
96
97
98
99
100
101
102
103
104
105
106
107
        # test bias
        model = DeformConv2d(1, 1, 2, stride=1, padding=0)
        assert not hasattr(model, 'bias')
        # test bias=True
        with pytest.raises(AssertionError):
            model = DeformConv2d(1, 1, 2, stride=1, padding=0, bias=True)
        # test in_channels % group != 0
        with pytest.raises(AssertionError):
            model = DeformConv2d(3, 2, 3, groups=2)
        # test out_channels % group != 0
        with pytest.raises(AssertionError):
            model = DeformConv2d(3, 4, 3, groups=3)

108
109
110
111
112
    def _test_amp_deformconv(self,
                             input_dtype,
                             threshold=1e-3,
                             batch_size=10,
                             im2col_step=2):
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        """The function to test amp released on pytorch 1.6.0.

        The type of input data might be torch.float or torch.half,
        so we should test deform_conv in both cases. With amp, the
        data type of model will NOT be set manually.

        Args:
            input_dtype: torch.float or torch.half.
            threshold: the same as above function.
        """
        if not torch.cuda.is_available():
            return
        from mmcv.ops import DeformConv2dPack
        c_in = 1
        c_out = 1
128
129
130
131
        repeated_input = np.repeat(input, batch_size, axis=0)
        repeated_gt_out = np.repeat(gt_out, batch_size, axis=0)
        repeated_gt_x_grad = np.repeat(gt_x_grad, batch_size, axis=0)
        x = torch.Tensor(repeated_input).cuda().type(input_dtype)
132
        x.requires_grad = True
133
134
135
136
137
138
139
        model = DeformConv2dPack(
            in_channels=c_in,
            out_channels=c_out,
            kernel_size=2,
            stride=1,
            padding=0,
            im2col_step=im2col_step)
140
141
142
143
144
145
146
147
148
149
150
        model.conv_offset.weight.data = torch.nn.Parameter(
            torch.Tensor(offset_weight).reshape(8, 1, 2, 2))
        model.conv_offset.bias.data = torch.nn.Parameter(
            torch.Tensor(offset_bias).reshape(8))
        model.weight.data = torch.nn.Parameter(
            torch.Tensor(deform_weight).reshape(1, 1, 2, 2))
        model.cuda()

        out = model(x)
        out.backward(torch.ones_like(out))

151
152
153
154
        assert np.allclose(out.data.detach().cpu().numpy(), repeated_gt_out,
                           threshold)
        assert np.allclose(x.grad.detach().cpu().numpy(), repeated_gt_x_grad,
                           threshold)
155
        assert np.allclose(
156
            model.conv_offset.weight.grad.detach().cpu().numpy() / batch_size,
157
            gt_offset_weight_grad, threshold)
158
159
160
161
162
163
        assert np.allclose(
            model.conv_offset.bias.grad.detach().cpu().numpy() / batch_size,
            gt_offset_bias_grad, threshold)
        assert np.allclose(
            model.weight.grad.detach().cpu().numpy() / batch_size,
            gt_deform_weight_grad, threshold)
164
165

        from mmcv.ops import DeformConv2d
166

167
168
169
170
171
172
173
174
175
176
177
178
179
        # test bias
        model = DeformConv2d(1, 1, 2, stride=1, padding=0)
        assert not hasattr(model, 'bias')
        # test bias=True
        with pytest.raises(AssertionError):
            model = DeformConv2d(1, 1, 2, stride=1, padding=0, bias=True)
        # test in_channels % group != 0
        with pytest.raises(AssertionError):
            model = DeformConv2d(3, 2, 3, groups=2)
        # test out_channels % group != 0
        with pytest.raises(AssertionError):
            model = DeformConv2d(3, 4, 3, groups=3)

180
    def test_deformconv(self):
181
182
        self._test_deformconv(torch.double, device='cpu')
        self._test_deformconv(torch.float, device='cpu', threshold=1e-1)
183
184
        self._test_deformconv(torch.double)
        self._test_deformconv(torch.float)
185
        self._test_deformconv(torch.half, threshold=1e-1)
186
187
188
189
190
191
192
        # test batch_size < im2col_step
        self._test_deformconv(torch.float, batch_size=1, im2col_step=2)
        # test bach_size % im2col_step != 0
        with pytest.raises(
                AssertionError,
                match='batch size must be divisible by im2col_step'):
            self._test_deformconv(torch.float, batch_size=10, im2col_step=3)
193
194
195

        # test amp when torch version >= '1.6.0', the type of
        # input data for deformconv might be torch.float or torch.half
196
        if (TORCH_VERSION != 'parrots'
197
                and digit_version(TORCH_VERSION) >= digit_version('1.6.0')):
198
199
200
            with autocast(enabled=True):
                self._test_amp_deformconv(torch.float, 1e-1)
                self._test_amp_deformconv(torch.half, 1e-1)