def test_reconstruction(self):
weights = (5, 5, 10, 10)
b2b_tfm = Box2BoxTransform(weights=weights)
src_boxes = random_boxes(10)
dst_boxes = random_boxes(10)
devices = [torch.device("cpu")]
if torch.cuda.is_available():
devices.append(torch.device("cuda"))
for device in devices:
src_boxes = src_boxes.to(device=device)
dst_boxes = dst_boxes.to(device=device)
deltas = b2b_tfm.get_deltas(src_boxes, dst_boxes)
dst_boxes_reconstructed = b2b_tfm.apply_deltas(deltas, src_boxes)
self.assertTrue(torch.allclose(dst_boxes, dst_boxes_reconstructed))
def _test_scriptability(self, device):
pooler_resolution = 14
canonical_level = 4
canonical_scale_factor = 2**canonical_level
pooler_scales = (1.0 / canonical_scale_factor, )
sampling_ratio = 0
N, C, H, W = 2, 4, 10, 8
N_rois = 10
std = 11
mean = 0
feature = (torch.rand(N, C, H, W) - 0.5) * 2 * std + mean
features = [feature.to(device)]
rois = []
for _ in range(N):
boxes = random_boxes(N_rois, W * canonical_scale_factor)
rois.append(Boxes(boxes).to(device))
roialignv2_pooler = ROIPooler(
output_size=pooler_resolution,
scales=pooler_scales,
sampling_ratio=sampling_ratio,
pooler_type="ROIAlignV2",
)
roialignv2_out = roialignv2_pooler(features, rois)
scripted_roialignv2_out = torch.jit.script(roialignv2_pooler)(features,
rois)
self.assertTrue(torch.equal(roialignv2_out, scripted_roialignv2_out))
def test_roi_pooler_tracing(self):
class Model(torch.nn.Module):
def __init__(self, roi):
super(Model, self).__init__()
self.roi = roi
def forward(self, x, boxes):
return self.roi(x, [Boxes(boxes)])
pooler_resolution = 14
canonical_level = 4
canonical_scale_factor = 2**canonical_level
pooler_scales = (1.0 / canonical_scale_factor,
0.5 / canonical_scale_factor)
sampling_ratio = 0
N, C, H, W = 1, 4, 10, 8
N_rois = 10
std = 11
mean = 0
feature = (torch.rand(N, C, H, W) - 0.5) * 2 * std + mean
feature = [feature, feature]
rois = random_boxes(N_rois, W * canonical_scale_factor)
# Add one larger box so that this level has only one box.
# This may trigger the bug https://github.com/pytorch/pytorch/issues/49852
# that we shall workaround.
rois = torch.cat([rois, torch.tensor([[0, 0, 448, 448]])])
model = Model(
ROIPooler(
output_size=pooler_resolution,
scales=pooler_scales,
sampling_ratio=sampling_ratio,
pooler_type="ROIAlign",
))
with torch.no_grad():
func = torch.jit.trace(model, (feature, rois))
o = func(feature, rois)
self.assertEqual(o.shape, (11, 4, 14, 14))
o = func(feature, rois[:5])
self.assertEqual(o.shape, (5, 4, 14, 14))
o = func(feature, random_boxes(20, W * canonical_scale_factor))
self.assertEqual(o.shape, (20, 4, 14, 14))
def test_paste_mask_scriptable(self):
scripted_f = torch.jit.script(paste_masks_in_image)
N = 10
masks = torch.rand(N, 28, 28)
boxes = Boxes(random_boxes(N, 100)).tensor
image_shape = (150, 150)
out = paste_masks_in_image(masks, boxes, image_shape)
scripted_out = scripted_f(masks, boxes, image_shape)
self.assertTrue(torch.equal(out, scripted_out))
def test_roi_pooler_tracing(self):
class Model(torch.nn.Module):
def __init__(self, roi):
super(Model, self).__init__()
self.roi = roi
def forward(self, x, boxes):
return self.roi(x, [Boxes(boxes)])
pooler_resolution = 14
canonical_level = 4
canonical_scale_factor = 2**canonical_level
pooler_scales = (1.0 / canonical_scale_factor,
0.5 / canonical_scale_factor)
sampling_ratio = 0
N, C, H, W = 1, 4, 10, 8
N_rois = 10
std = 11
mean = 0
feature = (torch.rand(N, C, H, W) - 0.5) * 2 * std + mean
feature = [feature, feature]
rois = random_boxes(N_rois, W * canonical_scale_factor)
model = Model(
ROIPooler(
output_size=pooler_resolution,
scales=pooler_scales,
sampling_ratio=sampling_ratio,
pooler_type="ROIAlign",
))
with torch.no_grad():
func = torch.jit.trace(model, (feature, rois))
o = func(feature, rois)
self.assertEqual(o.shape, (10, 4, 14, 14))
o = func(feature, rois[:5])
self.assertEqual(o.shape, (5, 4, 14, 14))
o = func(feature, random_boxes(20, W * canonical_scale_factor))
self.assertEqual(o.shape, (20, 4, 14, 14))
def test_reconstruction(self):
b2b_tfm = Box2BoxTransformLinear()
src_boxes = random_boxes(10)
dst_boxes = torch.tensor([0, 0, 101, 101] * 10).reshape(10, 4).float()
devices = [torch.device("cpu")]
if torch.cuda.is_available():
devices.append(torch.device("cuda"))
for device in devices:
src_boxes = src_boxes.to(device=device)
dst_boxes = dst_boxes.to(device=device)
deltas = b2b_tfm.get_deltas(src_boxes, dst_boxes)
dst_boxes_reconstructed = b2b_tfm.apply_deltas(deltas, src_boxes)
self.assertTrue(
torch.allclose(dst_boxes, dst_boxes_reconstructed, atol=1e-3))
def _test_roialignv2_roialignrotated_match(self, device):
pooler_resolution = 14
canonical_level = 4
canonical_scale_factor = 2**canonical_level
pooler_scales = (1.0 / canonical_scale_factor, )
sampling_ratio = 0
N, C, H, W = 2, 4, 10, 8
N_rois = 10
std = 11
mean = 0
feature = (torch.rand(N, C, H, W) - 0.5) * 2 * std + mean
features = [feature.to(device)]
rois = []
rois_rotated = []
for _ in range(N):
boxes = random_boxes(N_rois, W * canonical_scale_factor)
rotated_boxes = torch.zeros(N_rois, 5)
rotated_boxes[:, 0] = (boxes[:, 0] + boxes[:, 2]) / 2.0
rotated_boxes[:, 1] = (boxes[:, 1] + boxes[:, 3]) / 2.0
rotated_boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
rotated_boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
rois.append(Boxes(boxes).to(device))
rois_rotated.append(RotatedBoxes(rotated_boxes).to(device))
roialignv2_pooler = ROIPooler(
output_size=pooler_resolution,
scales=pooler_scales,
sampling_ratio=sampling_ratio,
pooler_type="ROIAlignV2",
)
roialignv2_out = roialignv2_pooler(features, rois)
roialignrotated_pooler = ROIPooler(
output_size=pooler_resolution,
scales=pooler_scales,
sampling_ratio=sampling_ratio,
pooler_type="ROIAlignRotated",
)
roialignrotated_out = roialignrotated_pooler(features, rois_rotated)
self.assertTrue(
torch.allclose(roialignv2_out, roialignrotated_out, atol=1e-4))
def gen_input(num1, num2):
feat = torch.randn((num1 + num2, 4, 14, 14))
box1 = random_boxes(num1)
box2 = random_boxes(num2)
return feat, box1, box2
def _create_tensors(self, N, device="cpu"):
boxes = random_boxes(N, 200, device=device)
scores = torch.rand(N, device=device)
return boxes, scores
def gen_inputs(h, w, N):
p2 = torch.rand(1, chan, h, w)
p3 = torch.rand(1, chan, h // 2, w // 2)
boxes = random_boxes(N, max_coord=h)
return p2, p3, boxes
def _create_tensors(self, N):
boxes = random_boxes(N, 200)
scores = torch.rand(N)
return boxes, scores
