mirror of https://github.com/FoundationVision/GLEE
172 lines
6.9 KiB
Python
172 lines
6.9 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
import logging
|
|
import unittest
|
|
import torch
|
|
|
|
from detectron2.layers import ShapeSpec
|
|
from detectron2.modeling.box_regression import Box2BoxTransform, Box2BoxTransformRotated
|
|
from detectron2.modeling.roi_heads.fast_rcnn import FastRCNNOutputLayers
|
|
from detectron2.modeling.roi_heads.rotated_fast_rcnn import RotatedFastRCNNOutputLayers
|
|
from detectron2.structures import Boxes, Instances, RotatedBoxes
|
|
from detectron2.utils.events import EventStorage
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FastRCNNTest(unittest.TestCase):
|
|
def test_fast_rcnn(self):
|
|
torch.manual_seed(132)
|
|
|
|
box_head_output_size = 8
|
|
|
|
box_predictor = FastRCNNOutputLayers(
|
|
ShapeSpec(channels=box_head_output_size),
|
|
box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)),
|
|
num_classes=5,
|
|
)
|
|
feature_pooled = torch.rand(2, box_head_output_size)
|
|
predictions = box_predictor(feature_pooled)
|
|
|
|
proposal_boxes = torch.tensor([[0.8, 1.1, 3.2, 2.8], [2.3, 2.5, 7, 8]], dtype=torch.float32)
|
|
gt_boxes = torch.tensor([[1, 1, 3, 3], [2, 2, 6, 6]], dtype=torch.float32)
|
|
proposal = Instances((10, 10))
|
|
proposal.proposal_boxes = Boxes(proposal_boxes)
|
|
proposal.gt_boxes = Boxes(gt_boxes)
|
|
proposal.gt_classes = torch.tensor([1, 2])
|
|
|
|
with EventStorage(): # capture events in a new storage to discard them
|
|
losses = box_predictor.losses(predictions, [proposal])
|
|
|
|
expected_losses = {
|
|
"loss_cls": torch.tensor(1.7951188087),
|
|
"loss_box_reg": torch.tensor(4.0357131958),
|
|
}
|
|
for name in expected_losses.keys():
|
|
assert torch.allclose(losses[name], expected_losses[name])
|
|
|
|
def test_fast_rcnn_empty_batch(self, device="cpu"):
|
|
box_predictor = FastRCNNOutputLayers(
|
|
ShapeSpec(channels=10),
|
|
box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)),
|
|
num_classes=8,
|
|
).to(device=device)
|
|
|
|
logits = torch.randn(0, 100, requires_grad=True, device=device)
|
|
deltas = torch.randn(0, 4, requires_grad=True, device=device)
|
|
losses = box_predictor.losses([logits, deltas], [])
|
|
for value in losses.values():
|
|
self.assertTrue(torch.allclose(value, torch.zeros_like(value)))
|
|
sum(losses.values()).backward()
|
|
self.assertTrue(logits.grad is not None)
|
|
self.assertTrue(deltas.grad is not None)
|
|
|
|
predictions, _ = box_predictor.inference([logits, deltas], [])
|
|
self.assertEqual(len(predictions), 0)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
|
|
def test_fast_rcnn_empty_batch_cuda(self):
|
|
self.test_fast_rcnn_empty_batch(device=torch.device("cuda"))
|
|
|
|
def test_fast_rcnn_rotated(self):
|
|
torch.manual_seed(132)
|
|
box_head_output_size = 8
|
|
|
|
box_predictor = RotatedFastRCNNOutputLayers(
|
|
ShapeSpec(channels=box_head_output_size),
|
|
box2box_transform=Box2BoxTransformRotated(weights=(10, 10, 5, 5, 1)),
|
|
num_classes=5,
|
|
)
|
|
feature_pooled = torch.rand(2, box_head_output_size)
|
|
predictions = box_predictor(feature_pooled)
|
|
proposal_boxes = torch.tensor(
|
|
[[2, 1.95, 2.4, 1.7, 0], [4.65, 5.25, 4.7, 5.5, 0]], dtype=torch.float32
|
|
)
|
|
gt_boxes = torch.tensor([[2, 2, 2, 2, 0], [4, 4, 4, 4, 0]], dtype=torch.float32)
|
|
proposal = Instances((10, 10))
|
|
proposal.proposal_boxes = RotatedBoxes(proposal_boxes)
|
|
proposal.gt_boxes = RotatedBoxes(gt_boxes)
|
|
proposal.gt_classes = torch.tensor([1, 2])
|
|
|
|
with EventStorage(): # capture events in a new storage to discard them
|
|
losses = box_predictor.losses(predictions, [proposal])
|
|
|
|
# Note: the expected losses are slightly different even if
|
|
# the boxes are essentially the same as in the FastRCNNOutput test, because
|
|
# bbox_pred in FastRCNNOutputLayers have different Linear layers/initialization
|
|
# between the two cases.
|
|
expected_losses = {
|
|
"loss_cls": torch.tensor(1.7920907736),
|
|
"loss_box_reg": torch.tensor(4.0410838127),
|
|
}
|
|
for name in expected_losses.keys():
|
|
assert torch.allclose(losses[name], expected_losses[name])
|
|
|
|
def test_predict_boxes_tracing(self):
|
|
class Model(torch.nn.Module):
|
|
def __init__(self, output_layer):
|
|
super(Model, self).__init__()
|
|
self._output_layer = output_layer
|
|
|
|
def forward(self, proposal_deltas, proposal_boxes):
|
|
instances = Instances((10, 10))
|
|
instances.proposal_boxes = Boxes(proposal_boxes)
|
|
return self._output_layer.predict_boxes((None, proposal_deltas), [instances])
|
|
|
|
box_head_output_size = 8
|
|
|
|
box_predictor = FastRCNNOutputLayers(
|
|
ShapeSpec(channels=box_head_output_size),
|
|
box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)),
|
|
num_classes=5,
|
|
)
|
|
|
|
model = Model(box_predictor)
|
|
|
|
from detectron2.export.torchscript_patch import patch_builtin_len
|
|
|
|
with torch.no_grad(), patch_builtin_len():
|
|
func = torch.jit.trace(model, (torch.randn(10, 20), torch.randn(10, 4)))
|
|
|
|
o = func(torch.randn(10, 20), torch.randn(10, 4))
|
|
self.assertEqual(o[0].shape, (10, 20))
|
|
o = func(torch.randn(5, 20), torch.randn(5, 4))
|
|
self.assertEqual(o[0].shape, (5, 20))
|
|
o = func(torch.randn(20, 20), torch.randn(20, 4))
|
|
self.assertEqual(o[0].shape, (20, 20))
|
|
|
|
def test_predict_probs_tracing(self):
|
|
class Model(torch.nn.Module):
|
|
def __init__(self, output_layer):
|
|
super(Model, self).__init__()
|
|
self._output_layer = output_layer
|
|
|
|
def forward(self, scores, proposal_boxes):
|
|
instances = Instances((10, 10))
|
|
instances.proposal_boxes = Boxes(proposal_boxes)
|
|
return self._output_layer.predict_probs((scores, None), [instances])
|
|
|
|
box_head_output_size = 8
|
|
|
|
box_predictor = FastRCNNOutputLayers(
|
|
ShapeSpec(channels=box_head_output_size),
|
|
box2box_transform=Box2BoxTransform(weights=(10, 10, 5, 5)),
|
|
num_classes=5,
|
|
)
|
|
|
|
model = Model(box_predictor)
|
|
|
|
from detectron2.export.torchscript_patch import patch_builtin_len
|
|
|
|
with torch.no_grad(), patch_builtin_len():
|
|
func = torch.jit.trace(model, (torch.randn(10, 6), torch.rand(10, 4)))
|
|
o = func(torch.randn(10, 6), torch.randn(10, 4))
|
|
self.assertEqual(o[0].shape, (10, 6))
|
|
o = func(torch.randn(5, 6), torch.randn(5, 4))
|
|
self.assertEqual(o[0].shape, (5, 6))
|
|
o = func(torch.randn(20, 6), torch.randn(20, 4))
|
|
self.assertEqual(o[0].shape, (20, 6))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|