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_disable_warnings(self) -> None:
"""
Tests .unsupported_ops_warnings(...) and .tracer_warnings(...)
"""
model = TraceWarningNet()
inputs = (torch.randn((1, *model.input_size)),)
analyzer = FlopCountAnalysis(model=model, inputs=inputs)
# Tracer warnings
analyzer.tracer_warnings(mode="all")
analyzer._stats = None # Manually clear cache so trace is rerun
self.assertWarns(torch.jit._trace.TracerWarning, analyzer.total)
analyzer._stats = None # Manually clear cache so trace is rerun
self.assertWarns(RuntimeWarning, analyzer.total)
analyzer.tracer_warnings(mode="none")
analyzer._stats = None # Manually clear cache so trace is rerun
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
_ = analyzer.total()
if w:
warning_types = [s.category for s in w]
self.assertFalse(torch.jit._trace.TracerWarning in warning_types)
self.assertFalse(RuntimeWarning in warning_types)
analyzer.tracer_warnings(mode="no_tracer_warning")
analyzer._stats = None # Manually clear cache so trace is rerun
self.assertWarns(RuntimeWarning, analyzer.total)
analyzer._stats = None # Manually clear cache so trace is rerun
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter("always")
_ = analyzer.total()
if w:
warning_types = [s.category for s in w]
self.assertFalse(torch.jit._trace.TracerWarning in warning_types)
# Unsupported ops and uncalled modules warnings
logger = logging.getLogger()
skipeed_msg = "Unsupported operator aten::add encountered 1 time(s)"
uncalled_msg = "never called"
uncalled_modules = "fc1" # fc2 is called by chance
analyzer.uncalled_modules_warnings(enabled=False)
analyzer.unsupported_ops_warnings(enabled=False)
analyzer._stats = None # Manually clear cache so trace is rerun
with self.assertLogs(logger, logging.WARN) as cm:
logger.warning("Dummy warning.")
_ = analyzer.total()
self.assertFalse(any(skipeed_msg in s for s in cm.output))
self.assertFalse(any(uncalled_msg in s for s in cm.output))
analyzer.unsupported_ops_warnings(enabled=True)
analyzer.uncalled_modules_warnings(enabled=True)
analyzer._stats = None # Manually clear cache so trace is rerun
with self.assertLogs(logger, logging.WARN) as cm:
_ = analyzer.total()
self.assertTrue(any(skipeed_msg in s for s in cm.output))
self.assertTrue(any(uncalled_msg in s for s in cm.output))
self.assertTrue(any(uncalled_modules in s for s in cm.output))
def test_by_module_and_operator(self) -> None:
"""
Tests that JitModelAnalysis.by_module_and_operator() returns
the correct counts in the correct structure.
"""
model = NestedNet(lin_op=self.lin_op)
inputs = (torch.randn((1, *model.input_size)), )
analyzer = FlopCountAnalysis(model=model, inputs=inputs)
analyzer.unsupported_ops_warnings(enabled=False)
self.assertEqual(analyzer.by_module_and_operator(), model.flops)
def test_by_module(self) -> None:
"""
Tests that JitModelAnalysis.by_module() returns the correct
counts in the correctly structured dictionary.
"""
model = NestedNet(lin_op=self.lin_op)
inputs = (torch.randn((1, *model.input_size)),)
analyzer = FlopCountAnalysis(model=model, inputs=inputs)
analyzer.unsupported_ops_warnings(enabled=False)
flops = {name: sum(counts.values()) for name, counts in model.flops.items()}
self.assertEqual(analyzer.by_module(), flops)
def test_by_operator(self) -> None:
"""
Tests that JitModelAnalysis.by_operator(module) returns the correct
counts for string and module inputs.
"""
model = NestedNet(lin_op=self.lin_op)
inputs = (torch.randn((1, *model.input_size)), )
analyzer = FlopCountAnalysis(model=model, inputs=inputs)
analyzer.unsupported_ops_warnings(enabled=False)
# Using a string input
for name in model.flops:
with self.subTest(name=name):
self.assertEqual(analyzer.by_operator(name), model.flops[name])
def test_skip_uncalled_containers_warnings(self) -> None:
# uncalled containers should not warn
class A(nn.Module):
def forward(self, x):
return self.submod[0](x) + 1
mod = A()
mod.submod = nn.ModuleList([nn.Linear(3, 3)]) # pyre-ignore
analyzer = FlopCountAnalysis(model=mod, inputs=torch.rand(1, 3))
analyzer.unsupported_ops_warnings(enabled=False)
logger = logging.getLogger()
with self.assertLogs(logger, logging.WARN) as cm:
logger.warning("Dummy warning.")
_ = analyzer.total()
uncalled_string = "Module never called: submod"
self.assertFalse(any(uncalled_string in s for s in cm.output))