mirror of https://github.com/FoundationVision/GLEE
211 lines
7.6 KiB
Python
211 lines
7.6 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
import numpy as np
|
|
import unittest
|
|
from copy import copy
|
|
import cv2
|
|
import torch
|
|
from fvcore.common.benchmark import benchmark
|
|
from torch.nn import functional as F
|
|
|
|
from detectron2.layers.roi_align import ROIAlign, roi_align
|
|
|
|
|
|
class ROIAlignTest(unittest.TestCase):
|
|
def test_forward_output(self):
|
|
input = np.arange(25).reshape(5, 5).astype("float32")
|
|
"""
|
|
0 1 2 3 4
|
|
5 6 7 8 9
|
|
10 11 12 13 14
|
|
15 16 17 18 19
|
|
20 21 22 23 24
|
|
"""
|
|
|
|
output = self._simple_roialign(input, [1, 1, 3, 3], (4, 4), aligned=False)
|
|
output_correct = self._simple_roialign(input, [1, 1, 3, 3], (4, 4), aligned=True)
|
|
|
|
# without correction:
|
|
old_results = [
|
|
[7.5, 8, 8.5, 9],
|
|
[10, 10.5, 11, 11.5],
|
|
[12.5, 13, 13.5, 14],
|
|
[15, 15.5, 16, 16.5],
|
|
]
|
|
|
|
# with 0.5 correction:
|
|
correct_results = [
|
|
[4.5, 5.0, 5.5, 6.0],
|
|
[7.0, 7.5, 8.0, 8.5],
|
|
[9.5, 10.0, 10.5, 11.0],
|
|
[12.0, 12.5, 13.0, 13.5],
|
|
]
|
|
# This is an upsampled version of [[6, 7], [11, 12]]
|
|
|
|
self.assertTrue(np.allclose(output.flatten(), np.asarray(old_results).flatten()))
|
|
self.assertTrue(
|
|
np.allclose(output_correct.flatten(), np.asarray(correct_results).flatten())
|
|
)
|
|
|
|
# Also see similar issues in tensorflow at
|
|
# https://github.com/tensorflow/tensorflow/issues/26278
|
|
|
|
def test_resize(self):
|
|
H, W = 30, 30
|
|
input = np.random.rand(H, W).astype("float32") * 100
|
|
box = [10, 10, 20, 20]
|
|
output = self._simple_roialign(input, box, (5, 5), aligned=True)
|
|
|
|
input2x = cv2.resize(input, (W // 2, H // 2), interpolation=cv2.INTER_LINEAR)
|
|
box2x = [x / 2 for x in box]
|
|
output2x = self._simple_roialign(input2x, box2x, (5, 5), aligned=True)
|
|
diff = np.abs(output2x - output)
|
|
self.assertTrue(diff.max() < 1e-4)
|
|
|
|
def test_grid_sample_equivalence(self):
|
|
H, W = 30, 30
|
|
input = np.random.rand(H, W).astype("float32") * 100
|
|
box = [10, 10, 20, 20]
|
|
for ratio in [1, 2, 3]:
|
|
output = self._simple_roialign(input, box, (5, 5), sampling_ratio=ratio)
|
|
output_grid_sample = grid_sample_roi_align(
|
|
torch.from_numpy(input[None, None, :, :]).float(),
|
|
torch.as_tensor(box).float()[None, :],
|
|
5,
|
|
1.0,
|
|
ratio,
|
|
)
|
|
self.assertTrue(torch.allclose(output, output_grid_sample))
|
|
|
|
def _simple_roialign(self, img, box, resolution, sampling_ratio=0, aligned=True):
|
|
"""
|
|
RoiAlign with scale 1.0.
|
|
"""
|
|
if isinstance(resolution, int):
|
|
resolution = (resolution, resolution)
|
|
op = ROIAlign(resolution, 1.0, sampling_ratio, aligned=aligned)
|
|
input = torch.from_numpy(img[None, None, :, :].astype("float32"))
|
|
|
|
rois = [0] + list(box)
|
|
rois = torch.from_numpy(np.asarray(rois)[None, :].astype("float32"))
|
|
output = op.forward(input, rois)
|
|
if torch.cuda.is_available():
|
|
output_cuda = op.forward(input.cuda(), rois.cuda()).cpu()
|
|
self.assertTrue(torch.allclose(output, output_cuda))
|
|
return output[0, 0]
|
|
|
|
def _simple_roialign_with_grad(self, img, box, resolution, device):
|
|
if isinstance(resolution, int):
|
|
resolution = (resolution, resolution)
|
|
|
|
op = ROIAlign(resolution, 1.0, 0, aligned=True)
|
|
input = torch.from_numpy(img[None, None, :, :].astype("float32"))
|
|
|
|
rois = [0] + list(box)
|
|
rois = torch.from_numpy(np.asarray(rois)[None, :].astype("float32"))
|
|
input = input.to(device=device)
|
|
rois = rois.to(device=device)
|
|
input.requires_grad = True
|
|
output = op.forward(input, rois)
|
|
return input, output
|
|
|
|
def test_empty_box(self):
|
|
img = np.random.rand(5, 5)
|
|
box = [3, 4, 5, 4]
|
|
o = self._simple_roialign(img, box, 7)
|
|
self.assertTrue(o.shape == (7, 7))
|
|
self.assertTrue((o == 0).all())
|
|
|
|
for dev in ["cpu"] + ["cuda"] if torch.cuda.is_available() else []:
|
|
input, output = self._simple_roialign_with_grad(img, box, 7, torch.device(dev))
|
|
output.sum().backward()
|
|
self.assertTrue(torch.allclose(input.grad, torch.zeros_like(input)))
|
|
|
|
def test_empty_batch(self):
|
|
input = torch.zeros(0, 3, 10, 10, dtype=torch.float32)
|
|
rois = torch.zeros(0, 5, dtype=torch.float32)
|
|
op = ROIAlign((7, 7), 1.0, 0, aligned=True)
|
|
output = op.forward(input, rois)
|
|
self.assertTrue(output.shape == (0, 3, 7, 7))
|
|
|
|
|
|
def grid_sample_roi_align(input, boxes, output_size, scale, sampling_ratio):
|
|
# unlike true roi_align, this does not support different batch_idx
|
|
from detectron2.projects.point_rend.point_features import (
|
|
generate_regular_grid_point_coords,
|
|
get_point_coords_wrt_image,
|
|
point_sample,
|
|
)
|
|
|
|
N, _, H, W = input.shape
|
|
R = len(boxes)
|
|
assert N == 1
|
|
boxes = boxes * scale
|
|
grid = generate_regular_grid_point_coords(R, output_size * sampling_ratio, device=boxes.device)
|
|
coords = get_point_coords_wrt_image(boxes, grid)
|
|
coords = coords / torch.as_tensor([W, H], device=coords.device) # R, s^2, 2
|
|
res = point_sample(input, coords.unsqueeze(0), align_corners=False) # 1,C, R,s^2
|
|
res = (
|
|
res.squeeze(0)
|
|
.permute(1, 0, 2)
|
|
.reshape(R, -1, output_size * sampling_ratio, output_size * sampling_ratio)
|
|
)
|
|
res = F.avg_pool2d(res, sampling_ratio)
|
|
return res
|
|
|
|
|
|
def benchmark_roi_align():
|
|
def random_boxes(mean_box, stdev, N, maxsize):
|
|
ret = torch.rand(N, 4) * stdev + torch.tensor(mean_box, dtype=torch.float)
|
|
ret.clamp_(min=0, max=maxsize)
|
|
return ret
|
|
|
|
def func(shape, nboxes_per_img, sampling_ratio, device, box_size="large"):
|
|
N, _, H, _ = shape
|
|
input = torch.rand(*shape)
|
|
boxes = []
|
|
batch_idx = []
|
|
for k in range(N):
|
|
if box_size == "large":
|
|
b = random_boxes([80, 80, 130, 130], 24, nboxes_per_img, H)
|
|
else:
|
|
b = random_boxes([100, 100, 110, 110], 4, nboxes_per_img, H)
|
|
boxes.append(b)
|
|
batch_idx.append(torch.zeros(nboxes_per_img, 1, dtype=torch.float32) + k)
|
|
boxes = torch.cat(boxes, axis=0)
|
|
batch_idx = torch.cat(batch_idx, axis=0)
|
|
boxes = torch.cat([batch_idx, boxes], axis=1)
|
|
|
|
input = input.to(device=device)
|
|
boxes = boxes.to(device=device)
|
|
|
|
def bench():
|
|
if False and sampling_ratio > 0 and N == 1:
|
|
# enable to benchmark grid_sample (slower)
|
|
grid_sample_roi_align(input, boxes[:, 1:], 7, 1.0, sampling_ratio)
|
|
else:
|
|
roi_align(input, boxes, 7, 1.0, sampling_ratio, True)
|
|
if device == "cuda":
|
|
torch.cuda.synchronize()
|
|
|
|
return bench
|
|
|
|
def gen_args(arg):
|
|
args = []
|
|
for size in ["small", "large"]:
|
|
for ratio in [0, 2]:
|
|
args.append(copy(arg))
|
|
args[-1]["sampling_ratio"] = ratio
|
|
args[-1]["box_size"] = size
|
|
return args
|
|
|
|
arg = dict(shape=(1, 512, 256, 256), nboxes_per_img=512, device="cuda")
|
|
benchmark(func, "cuda_roialign", gen_args(arg), num_iters=20, warmup_iters=1)
|
|
arg.update({"device": "cpu", "shape": (1, 256, 128, 128)})
|
|
benchmark(func, "cpu_roialign", gen_args(arg), num_iters=5, warmup_iters=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if torch.cuda.is_available():
|
|
benchmark_roi_align()
|
|
unittest.main()
|