mirror of https://github.com/open-mmlab/mmcv.git
92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
# 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)
|