mirror of https://github.com/FoundationVision/GLEE
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
# 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"})
|