def test_complexity_calculation(self) -> None:
model = TestModel()
input_shape = (3, 10, 10)
num_elems = 3 * 10 * 10
self.assertEqual(compute_activations(model, input_shape=input_shape),
num_elems)
self.assertEqual(
compute_flops(model, input_shape=input_shape),
num_elems + 0 +
(300 * 300), # TestModule + TestConvModule + TestModel.linear;
# TestModel.unused_linear is unused and shouldn't be counted
)
self.assertEqual(
count_params(model),
(2 * 3) + (2 * 3 * 4 * 4) + (4 * 5) + (300 * 300) + (10 * 10) +
(2 * 2),
) # TestModule.linear + TestConvModule + TestConvModule.linear +
# TestModel.linear + TestModel.extra_params + TestModel.unused_linear
# test that we calculate complexity correctly for a model which re-uses a module
model = TestModel2()
in_channels = 3
out_channels = 3
out_h, out_w = 9, 9
kernel_h, kernel_w = 2, 2
conv_flops = in_channels * out_channels * out_h * out_w * kernel_h * kernel_w
conv_activations = out_channels * out_h * out_w
self.assertEqual(compute_activations(model, input_shape=input_shape),
conv_activations * 2) # the conv is applied twice
self.assertEqual(compute_flops(model, input_shape=input_shape),
conv_flops * 2) # the conv is applied twice
self.assertEqual(count_params(model),
in_channels * out_channels * kernel_h * kernel_w)
def test_flops_calculation(self):
# test that a model containing a custom module which doesn't have FLOPs defined
# raises an exception
model = nn.Sequential(TestModuleWithoutFlops())
input_shape = (3, 10, 10)
with self.assertRaises(Exception):
compute_flops(model, input_shape=input_shape)
# test that a model containing a custom module does have FLOPs defined works,
# even if the module has children which don't define FLOPs
model = nn.Sequential(TestModuleWithFlops())
input_shape = (3, 10, 10)
self.assertEqual(
compute_flops(model, input_shape=input_shape),
TestModuleWithFlops._flops) # the conv is applied twice
def on_start(self, task) -> None:
"""
Before the training starts, run one forward only pass of the model on the
dummy input of shape specified by user in HOOKS.MODEL_COMPLEXITY.INPUT_SHAPE
We calculate the flops, activations and number of params in the model.
"""
self.num_flops, self.num_activations, self.num_parameters = 0, 0, 0
input_shape = task.config["HOOKS"]["MODEL_COMPLEXITY"]["INPUT_SHAPE"]
try:
self.num_parameters = count_params(task.base_model)
self.num_parameters = round(float(self.num_parameters) / 1000000, 4)
try:
self.num_flops = compute_flops(task.base_model, input_shape)
if self.num_flops is None:
logging.info("FLOPs for forward pass: skipped.")
self.num_flops = round(float(self.num_flops) / 1000000000, 4)
except NotImplementedError:
logging.warning(
"Unsupported modules found in model. FLOPs calculation skipped "
)
logging.debug("Exception: ", exc_info=True)
try:
self.num_activations = compute_activations(task.base_model, input_shape)
self.num_activations = round(float(self.num_activations) / 1000000, 4)
except NotImplementedError:
logging.info("input_shape not found. Skipping activation calculation.")
logging.info(
f"#params (10^6): {self.num_parameters} "
f"#FLOPs (10^9): {self.num_flops} "
f"#activations (10^6): {self.num_activations}"
)
except Exception:
logging.exception("Unexpected failure estimating model complexity.")
def on_start(self, task: "tasks.ClassyTask",
local_variables: Dict[str, Any]) -> None:
"""Measure number of parameters, FLOPs and activations."""
try:
num_flops = compute_flops(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key if hasattr(
task.base_model, "input_key") else None,
)
if num_flops is None:
logging.info("FLOPs for forward pass: skipped.")
else:
logging.info("FLOPs for forward pass: %d MFLOPs" %
(float(num_flops) / 1e6))
except NotImplementedError:
logging.warning(
"""Model contains unsupported modules:
Could not compute FLOPs for model forward pass. Exception:""",
exc_info=True,
)
num_activations = compute_activations(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key if hasattr(
task.base_model, "input_key") else None,
)
logging.info(f"Number of activations in model: {num_activations}")
logging.info("Number of parameters in model: %d" %
count_params(task.base_model))
def on_start(
self, task: "tasks.ClassyTask", local_variables: Dict[str, Any]
) -> None:
"""Measure number of parameters, FLOPs and activations."""
self.num_flops = 0
self.num_activations = 0
self.num_parameters = 0
try:
self.num_parameters = count_params(task.base_model)
logging.info("Number of parameters in model: %d" % self.num_parameters)
try:
self.num_flops = compute_flops(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key
if hasattr(task.base_model, "input_key")
else None,
)
if self.num_flops is None:
logging.info("FLOPs for forward pass: skipped.")
self.num_flops = 0
else:
logging.info(
"FLOPs for forward pass: %d MFLOPs"
% (float(self.num_flops) / 1e6)
)
except NotImplementedError:
logging.warning(
"""Model contains unsupported modules:
Could not compute FLOPs for model forward pass. Exception:""",
exc_info=True,
)
try:
self.num_activations = compute_activations(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key
if hasattr(task.base_model, "input_key")
else None,
)
logging.info(f"Number of activations in model: {self.num_activations}")
except NotImplementedError:
logging.info(
"Model does not implement input_shape. Skipping "
"activation calculation."
)
except Exception:
logging.exception("Unexpected failure estimating model complexity.")
def test_complexity_calculation_resnext(self) -> None:
model_configs = get_test_model_configs()
# make sure there are three configs returned
self.assertEqual(len(model_configs), 3)
# expected values which allow minor deviations from model changes
# we only test at the 10^6 scale
expected_m_flops = [4122, 7850, 8034]
expected_m_params = [25, 44, 44]
expected_m_activations = [11, 16, 21]
for model_config, m_flops, m_params, m_activations in zip(
model_configs, expected_m_flops, expected_m_params, expected_m_activations
):
model = build_model(model_config)
self.assertEqual(compute_activations(model) // 10 ** 6, m_activations)
self.assertEqual(compute_flops(model) // 10 ** 6, m_flops)
self.assertEqual(count_params(model) // 10 ** 6, m_params)
def on_start(self, task) -> None:
"""Measure number of parameters, FLOPs and activations."""
self.num_flops = 0
self.num_activations = 0
self.num_parameters = 0
try:
self.num_parameters = count_params(task.base_model)
logging.info("Number of parameters in model: %d" %
self.num_parameters)
try:
self.num_flops = compute_flops(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key if hasattr(
task.base_model, "input_key") else None,
)
if self.num_flops is None:
logging.info("FLOPs for forward pass: skipped.")
self.num_flops = 0
else:
logging.info("FLOPs for forward pass: %d MFLOPs" %
(float(self.num_flops) / 1e6))
except ClassyProfilerNotImplementedError as e:
logging.warning(
f"Could not compute FLOPs for model forward pass: {e}")
try:
self.num_activations = compute_activations(
task.base_model,
input_shape=task.base_model.input_shape,
input_key=task.base_model.input_key if hasattr(
task.base_model, "input_key") else None,
)
logging.info(
f"Number of activations in model: {self.num_activations}")
except ClassyProfilerNotImplementedError as e:
logging.warning(
f"Could not compute activations for model forward pass: {e}"
)
except Exception:
logging.info("Skipping complexity calculation: Unexpected error")
logging.debug("Error trace for complexity calculation:",
exc_info=True)
