def test_non_forward_func_call(self) -> None: """ Tests calls to a submodule's non-forward function. Also tests that the intermediate module is correctly identified as a skipped module. """ model = NonForwardNet() inputs = (torch.randn((1, 10)), ) analyzer = FlopCountAnalysis(model=model, inputs=inputs).ancestor_mode("caller") inner_fc_count = model.submod.fc_flops total_count = model.fc_flops + inner_fc_count self.assertEqual(analyzer.total("submod"), 0) self.assertEqual(analyzer.total("submod.fc"), inner_fc_count) self.assertEqual(analyzer.total(""), total_count) # The mod not directly called is registered as such self.assertEqual(analyzer.uncalled_modules(), {"submod"}) analyzer = FlopCountAnalysis(model=model, inputs=inputs).ancestor_mode("owner") self.assertEqual(analyzer.total("submod"), inner_fc_count) self.assertEqual(analyzer.total("submod.fc"), inner_fc_count) self.assertEqual(analyzer.total(""), total_count) self.assertEqual(analyzer.uncalled_modules(), set())
def test_data_parallel(self) -> None: """ Tests that a model wrapped in DataParallel still returns results labeled by the correct scopes. """ model = NestedNet(lin_op=self.lin_op) inputs = (torch.randn((1, *model.input_size)), ) # Find flops for wrapper flops = { "module" + ("." if name else "") + name: flop for name, flop in model.flops.items() } flops[""] = model.flops[""] name_to_module = { "module" + ("." if name else "") + name: mod for name, mod in model.name_to_module.items() } name_to_module[""] = model.name_to_module[""] model = torch.nn.DataParallel(model).cpu() analyzer = FlopCountAnalysis(model=model, inputs=inputs) analyzer.unsupported_ops_warnings(enabled=False) # Using a string input for name in flops: with self.subTest(name=name): gt_flops = sum(flops[name].values()) self.assertEqual(analyzer.total(name), gt_flops) # Output as dictionary self.assertEqual(analyzer.by_module_and_operator(), flops) # Test no uncalled modules self.assertEqual(analyzer.uncalled_modules(), set())
def test_shared_module(self) -> None: """ Tests the behavior of shared submodules that may have multiple names. """ model = SharedModuleNet() inputs = (torch.randn((1, *model.input_size)),) analyzer = FlopCountAnalysis( model=model, inputs=inputs ).unsupported_ops_warnings(enabled=False) # The names `submod2.submod` and `multiname2` are not included, # since only the first name of a module is made the canonical one. # The counts associated with these cases are included under # `submod1.submod` and `multiname1` respectively. multiname_flops = 2 * model.multiname_flops # Called under 2 names shared_flops = 2 * model.shared_flops # Shared under 2 submodules total_flops = multiname_flops + shared_flops flops = { "": total_flops, "submod1": model.shared_flops, "submod1.submod": shared_flops, "submod2": model.shared_flops, "multiname1": multiname_flops, } self.assertEqual(analyzer.by_module(), flops) # Test access by alternative name self.assertEqual( analyzer.total("submod2.submod"), flops["submod1.submod"], ) self.assertEqual( analyzer.total("multiname2"), flops["multiname1"], ) # Test getting canonical name self.assertEqual(analyzer.canonical_module_name("multiname2"), "multiname1") self.assertEqual(analyzer.canonical_module_name("multiname1"), "multiname1") self.assertEqual( analyzer.canonical_module_name("submod2.submod"), "submod1.submod" ) self.assertEqual( analyzer.canonical_module_name("submod1.submod"), "submod1.submod" ) # Tests no uncalled modules self.assertEqual(analyzer.uncalled_modules(), set())
def test_recursive_scope(self) -> None: """ Tests that an op is only counted once per module, even if it is in the scope of that module multiple times. """ model = RecursiveScopeNet() inputs = (torch.randn((1, *model.input_size)), ) analyzer = FlopCountAnalysis(model, inputs) self.assertEqual(analyzer.total(), model.flops) self.assertEqual(analyzer.total("fc"), model.flops) # Tests no uncalled modules self.assertEqual(analyzer.uncalled_modules(), set())
def test_repeated_module(self) -> None: """ Tests that repeated calls to the same submodule correct aggregates results to that submodule. """ model = RepeatedNet() inputs = (torch.randn((1, *model.input_size)), ) analyzer = FlopCountAnalysis(model=model, inputs=inputs) fc1_count = model.fc1_num * model.fc1_flops fc2_count = model.fc2_num * model.fc2_flops total_count = fc1_count + fc2_count fc1_per_operator = Counter({self.lin_op: fc1_count}) self.assertEqual(analyzer.total("fc1"), fc1_count) self.assertEqual(analyzer.total("fc2"), fc2_count) self.assertEqual(analyzer.total(""), total_count) self.assertEqual(analyzer.by_operator("fc1"), fc1_per_operator) # Tests no uncalled mods self.assertEqual(analyzer.uncalled_modules(), set())
def test_unused_module(self) -> None: """ Tests that unused modules return 0 count for operator sums and and empty Counter() for per-operator results. Also tests that unused modules are reported by .uncalled_modules(), but that modules that simply have zero flops (like ReLU) are not. """ model = UnusedNet() inputs = (torch.randn((1, *model.input_size)), ) analyzer = FlopCountAnalysis(model=model, inputs=inputs) unused_count = 0 unused_per_operator = Counter() model_count = model.fc1_flops + model.fc2_flops self.assertEqual(analyzer.total("unused"), unused_count) self.assertEqual(analyzer.by_operator("unused"), unused_per_operator) self.assertEqual(analyzer.total(""), model_count) # The unused mod is recognized as never called self.assertEqual(analyzer.uncalled_modules(), {"unused"})