mirror of https://github.com/open-mmlab/mmcv.git
158 lines
6.5 KiB
Python
158 lines
6.5 KiB
Python
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from mmcv.utils import TORCH_VERSION, digit_version
|
|
|
|
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
|
|
|
|
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):
|
|
|
|
def _test_deformconv(self,
|
|
dtype=torch.float,
|
|
threshold=1e-3,
|
|
device='cuda'):
|
|
if not torch.cuda.is_available() and device == 'cuda':
|
|
pytest.skip('test requires GPU')
|
|
from mmcv.ops import DeformConv2dPack
|
|
c_in = 1
|
|
c_out = 1
|
|
x = torch.tensor(input, device=device, dtype=dtype)
|
|
x.requires_grad = True
|
|
model = DeformConv2dPack(c_in, c_out, 2, stride=1, padding=0)
|
|
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))
|
|
if device == 'cuda':
|
|
model.cuda()
|
|
model.type(dtype)
|
|
|
|
out = model(x)
|
|
out.backward(torch.ones_like(out))
|
|
|
|
assert np.allclose(out.data.detach().cpu().numpy(), gt_out, threshold)
|
|
assert np.allclose(x.grad.detach().cpu().numpy(), gt_x_grad, threshold)
|
|
assert np.allclose(
|
|
model.conv_offset.weight.grad.detach().cpu().numpy(),
|
|
gt_offset_weight_grad, threshold)
|
|
assert np.allclose(model.conv_offset.bias.grad.detach().cpu().numpy(),
|
|
gt_offset_bias_grad, threshold)
|
|
assert np.allclose(model.weight.grad.detach().cpu().numpy(),
|
|
gt_deform_weight_grad, threshold)
|
|
|
|
from mmcv.ops import DeformConv2d
|
|
|
|
# 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)
|
|
|
|
def _test_amp_deformconv(self, input_dtype, threshold=1e-3):
|
|
"""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
|
|
x = torch.Tensor(input).cuda().type(input_dtype)
|
|
x.requires_grad = True
|
|
model = DeformConv2dPack(c_in, c_out, 2, stride=1, padding=0)
|
|
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))
|
|
|
|
assert np.allclose(out.data.detach().cpu().numpy(), gt_out, threshold)
|
|
assert np.allclose(x.grad.detach().cpu().numpy(), gt_x_grad, threshold)
|
|
assert np.allclose(
|
|
model.conv_offset.weight.grad.detach().cpu().numpy(),
|
|
gt_offset_weight_grad, threshold)
|
|
assert np.allclose(model.conv_offset.bias.grad.detach().cpu().numpy(),
|
|
gt_offset_bias_grad, threshold)
|
|
assert np.allclose(model.weight.grad.detach().cpu().numpy(),
|
|
gt_deform_weight_grad, threshold)
|
|
|
|
from mmcv.ops import DeformConv2d
|
|
|
|
# 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)
|
|
|
|
def test_deformconv(self):
|
|
self._test_deformconv(torch.double, device='cpu')
|
|
self._test_deformconv(torch.float, device='cpu', threshold=1e-1)
|
|
self._test_deformconv(torch.double)
|
|
self._test_deformconv(torch.float)
|
|
self._test_deformconv(torch.half, threshold=1e-1)
|
|
|
|
# test amp when torch version >= '1.6.0', the type of
|
|
# input data for deformconv might be torch.float or torch.half
|
|
if (TORCH_VERSION != 'parrots'
|
|
and digit_version(TORCH_VERSION) >= digit_version('1.6.0')):
|
|
with autocast(enabled=True):
|
|
self._test_amp_deformconv(torch.float, 1e-1)
|
|
self._test_amp_deformconv(torch.half, 1e-1)
|