mirror of https://github.com/open-mmlab/mmcv.git
75 lines
3.4 KiB
Python
75 lines
3.4 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from torch.autograd import gradcheck
|
|
|
|
from mmcv.ops import RiRoIAlignRotated
|
|
|
|
np_feature = np.array([[[[1, 2], [3, 4]], [[1, 2], [4, 3]], [[4, 3], [2, 1]],
|
|
[[1, 2], [5, 6]], [[3, 4], [7, 8]], [[9, 10], [13,
|
|
14]],
|
|
[[11, 12], [15, 16]], [[1, 1], [2, 2]]]])
|
|
np_rois = np.array([[0., 0.5, 0.5, 1., 1., np.pi / 3],
|
|
[0., 1., 1., 3., 3., np.pi / 2]])
|
|
expect_output = np.array([[[[1.8425, 1.3516], [2.3151, 1.8241]],
|
|
[[2.4779, 1.7416], [3.2173, 2.5632]],
|
|
[[2.7149, 2.2638], [2.6540, 2.3673]],
|
|
[[2.9461, 2.8638], [2.8028, 2.7205]],
|
|
[[4.1943, 2.7214], [5.6119, 4.1391]],
|
|
[[7.5276, 6.0547], [8.9453, 7.4724]],
|
|
[[12.1943, 10.7214], [13.6119, 12.1391]],
|
|
[[9.5489, 8.4237], [10.5763, 9.4511]]],
|
|
[[[7.6562, 12.5625], [4.0000, 6.6250]],
|
|
[[1.0000, 1.3125], [0.5000, 0.6562]],
|
|
[[1.6562, 1.9375], [1.0000, 1.3125]],
|
|
[[1.8438, 2.0547], [0.7500, 1.1562]],
|
|
[[0.8438, 3.0625], [0.2500, 1.1875]],
|
|
[[2.6562, 2.5625], [1.5000, 1.6250]],
|
|
[[3.6562, 4.5625], [2.0000, 2.6250]],
|
|
[[6.6562, 10.5625], [3.5000, 5.6250]]]])
|
|
|
|
expect_grad = np.array([[[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]],
|
|
[[1.4727, 1.5586], [1.5586, 1.6602]]]])
|
|
|
|
pool_h = 2
|
|
pool_w = 2
|
|
spatial_scale = 1.0
|
|
num_samples = 2
|
|
sampling_ratio = 2
|
|
num_orientations = 8
|
|
clockwise = False
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason='requires CUDA support')
|
|
def test_roialign_rotated_gradcheck():
|
|
x = torch.tensor(
|
|
np_feature, dtype=torch.float, device='cuda', requires_grad=True)
|
|
rois = torch.tensor(np_rois, dtype=torch.float, device='cuda')
|
|
froipool = RiRoIAlignRotated((pool_h, pool_w), spatial_scale, num_samples,
|
|
num_orientations, clockwise)
|
|
gradcheck(froipool, (x, rois), eps=1e-3, atol=1e-3)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason='requires CUDA support')
|
|
def test_roialign_rotated_allclose():
|
|
x = torch.tensor(
|
|
np_feature, dtype=torch.float, device='cuda', requires_grad=True)
|
|
rois = torch.tensor(np_rois, dtype=torch.float, device='cuda')
|
|
froipool = RiRoIAlignRotated((pool_h, pool_w), spatial_scale, num_samples,
|
|
num_orientations, clockwise)
|
|
output = froipool(x, rois)
|
|
output.backward(torch.ones_like(output))
|
|
assert np.allclose(
|
|
output.data.type(torch.float).cpu().numpy(), expect_output, atol=1e-3)
|
|
assert np.allclose(
|
|
x.grad.data.type(torch.float).cpu().numpy(), expect_grad, atol=1e-3)
|