Spaces:
Sleeping
Sleeping
File size: 3,601 Bytes
d7e58f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
# Copyright (c) OpenMMLab. All rights reserved.
import copy
import numpy as np
import pytest
import torch
# [1,4c,h,w]
input_arr = [[[[1., 2., 3., 4.], [5., 6., 7., 8.], [9., 10., 11., 12.]],
[[6, 7, 5, 8], [2, 1, 3, 4], [12, 9, 11, 10]],
[[-2, -3, 2, 0], [-4, -5, 1, -1], [-1, -1, -1, -1]],
[[0, -1, 2, 1], [-4, -3, -2, -1], [-1, -2, -3, -4]]]]
# [1,h*w,4]
boxes_arr = [[[0, 0, 2, 1], [1, 0, 3, 1], [1, 0, 2, 1], [0, 0, 3, 1],
[0, 0, 1, 2], [0, 0, 2, 2], [1, 0, 2, 1], [1, 0, 3, 1],
[0, 1, 1, 2], [0, 0, 3, 2], [1, 0, 3, 2], [2, 0, 3, 2]]]
output_dict = {
# [1,c,h*w,4] for each value,
# the output is manually checked for its correctness
# pool_size=1
1: [[[[3., 6., 1., 2.], [4., 7., -1., 1.], [3., 7., 1., 2.],
[4., 6., -1., 1.], [2., 12., -1., -1.], [3., 12., -1., 2.],
[3., 7., 1., 2.], [4., 7., -1., 1.], [6., 12., -1., -2.],
[4., 12., -1., 1.], [4., 9., -1., 1.], [4., 11., -1., 1.]]]],
# pool_size=2
2: [[[[3., 6., 1., 2.], [4., 7., 1., 1.], [3., 7., 1., 2.],
[4., 6., -1., 1.], [2., 12., -1., -1.], [3., 12., -1., 2.],
[3., 7., 1., 2.], [4., 7., 1., 1.], [6., 12., -1., -2.],
[4., 12., -1., 1.], [4., 9., -1., 1.], [4., 11., -1., 1.]]]],
}
input_grad_dict = {
# [1,4c,h,w] for each value
# the grad is manually checked for its correctness
# pool_size=1
1: [[[[0., 1., 4., 6.], [0., 1., 0., 0.], [0., 0., 0., 0.]],
[[2., 4., 0., 0.], [0., 0., 0., 0.], [4., 1., 1., 0.]],
[[0., 0., 0., 0.], [0., 0., 3., 3.], [0., 2., 1., 3.]],
[[0., 1., 4., 6.], [0., 0., 0., 0.], [0., 1., 0., 0.]]]],
# pool_size=2
2: [[[[0., 1., 4., 6.], [0., 1., 0., 0.], [0., 0., 0., 0.]],
[[2., 4., 0., 0.], [0., 0., 0., 0.], [4., 1., 1., 0.]],
[[0., 0., 0., 0.], [0., 0., 5., 1.], [0., 2., 1., 3.]],
[[0., 1., 4., 6.], [0., 0., 0., 0.], [0., 1., 0., 0.]]]],
}
def _test_border_align_allclose(device, dtype, pool_size):
if not torch.cuda.is_available() and device == 'cuda':
pytest.skip('test requires GPU')
try:
from mmcv.ops import BorderAlign, border_align
except ModuleNotFoundError:
pytest.skip('BorderAlign op is not successfully compiled')
np_input = np.array(input_arr)
np_boxes = np.array(boxes_arr)
np_output = np.array(output_dict[pool_size])
np_grad = np.array(input_grad_dict[pool_size])
input = torch.tensor(
np_input, dtype=dtype, device=device, requires_grad=True)
boxes = torch.tensor(np_boxes, dtype=dtype, device=device)
# test for border_align
input_cp = copy.deepcopy(input)
output = border_align(input_cp, boxes, pool_size)
output.backward(torch.ones_like(output))
assert np.allclose(
output.data.type(dtype).cpu().numpy(), np_output, atol=1e-5)
assert np.allclose(
input_cp.grad.data.type(dtype).cpu().numpy(), np_grad, atol=1e-5)
# test for BorderAlign
pool_module = BorderAlign(pool_size)
output = pool_module(input, boxes)
output.backward(torch.ones_like(output))
assert np.allclose(
output.data.type(dtype).cpu().numpy(), np_output, atol=1e-5)
assert np.allclose(
input.grad.data.type(dtype).cpu().numpy(), np_grad, atol=1e-5)
@pytest.mark.parametrize('device', ['cuda'])
@pytest.mark.parametrize('dtype', [torch.float, torch.half, torch.double])
@pytest.mark.parametrize('pool_size', [1, 2])
def test_border_align(device, dtype, pool_size):
_test_border_align_allclose(device, dtype, pool_size)
|