Spaces:
Running
Running
File size: 7,038 Bytes
6f6830f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
# 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()
|