Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| import unittest | |
| import torch | |
| from torch import nn | |
| from detectron2.utils.analysis import find_unused_parameters, flop_count_operators, parameter_count | |
| from detectron2.utils.testing import get_model_no_weights | |
| class RetinaNetTest(unittest.TestCase): | |
| def setUp(self): | |
| self.model = get_model_no_weights("COCO-Detection/retinanet_R_50_FPN_1x.yaml") | |
| def test_flop(self): | |
| # RetinaNet supports flop-counting with random inputs | |
| inputs = [{"image": torch.rand(3, 800, 800), "test_unused": "abcd"}] | |
| res = flop_count_operators(self.model, inputs) | |
| self.assertEqual(int(res["conv"]), 146) # 146B flops | |
| def test_param_count(self): | |
| res = parameter_count(self.model) | |
| self.assertEqual(res[""], 37915572) | |
| self.assertEqual(res["backbone"], 31452352) | |
| class FasterRCNNTest(unittest.TestCase): | |
| def setUp(self): | |
| self.model = get_model_no_weights("COCO-Detection/faster_rcnn_R_50_FPN_1x.yaml") | |
| def test_flop(self): | |
| # Faster R-CNN supports flop-counting with random inputs | |
| inputs = [{"image": torch.rand(3, 800, 800)}] | |
| res = flop_count_operators(self.model, inputs) | |
| # This only checks flops for backbone & proposal generator | |
| # Flops for box head is not conv, and depends on #proposals, which is | |
| # almost 0 for random inputs. | |
| self.assertEqual(int(res["conv"]), 117) | |
| def test_flop_with_output_shape(self): | |
| inputs = [{"image": torch.rand(3, 800, 800), "height": 700, "width": 700}] | |
| res = flop_count_operators(self.model, inputs) | |
| self.assertEqual(int(res["conv"]), 117) | |
| def test_param_count(self): | |
| res = parameter_count(self.model) | |
| self.assertEqual(res[""], 41699936) | |
| self.assertEqual(res["backbone"], 26799296) | |
| class MaskRCNNTest(unittest.TestCase): | |
| def setUp(self): | |
| self.model = get_model_no_weights("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml") | |
| def test_flop(self): | |
| inputs1 = [{"image": torch.rand(3, 800, 800)}] | |
| inputs2 = [{"image": torch.rand(3, 800, 800), "height": 700, "width": 700}] | |
| for inputs in [inputs1, inputs2]: | |
| res = flop_count_operators(self.model, inputs) | |
| # The mask head could have extra conv flops, so total >= 117 | |
| self.assertGreaterEqual(int(res["conv"]), 117) | |
| class UnusedParamTest(unittest.TestCase): | |
| def test_unused(self): | |
| class TestMod(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.fc1 = nn.Linear(10, 10) | |
| self.t = nn.Linear(10, 10) | |
| def forward(self, x): | |
| return self.fc1(x).mean() | |
| m = TestMod() | |
| ret = find_unused_parameters(m, torch.randn(10, 10)) | |
| self.assertEqual(set(ret), {"t.weight", "t.bias"}) | |