Add activation_count_operators to analysis
Summary: We are going to have support for activation count in fvcore. Adding the wrapper for D2 so that we can count activations in addition to FLOPs.
import logging
import typing
import torch
-from fvcore.nn import flop_count, parameter_count, parameter_count_table
+from fvcore.nn import activation_count, flop_count, parameter_count, parameter_count_table
from torch import nn
from .logger import log_first_n
-__all__ = ["flop_count_operators", "parameter_count_table", "parameter_count"]
+__all__ = [
+ "activation_count_operators",
+ "flop_count_operators",
+ "parameter_count_table",
+ "parameter_count",
+]
+
+FLOPS_MODE = "flops"
+ACTIVATIONS_MODE = "activations"
def flop_count_operators(
@@ -32,7 +40,35 @@ def flop_count_operators(
model: a detectron2 model that takes `list[dict]` as input.
inputs (list[dict]): inputs to model, in detectron2's standard format.
"""
- assert len(inputs) == 1, "Please only compute flops with batch size=1"
+ return _wrapper_count_operators(model=model, inputs=inputs, mode=FLOPS_MODE, **kwargs)
+
+
+def activation_count_operators(
+ model: nn.Module, inputs: list, **kwargs
+) -> typing.DefaultDict[str, float]:
+ """
+ Implement operator-level activations counting using jit.
+ This is a wrapper of fvcore.nn.activation_count, that supports standard detection models
+ in detectron2.
+
+ Note:
+ The function runs the input through the model to compute activations.
+ The activations of a detection model is often input-dependent, for example,
+ the activations of box & mask head depends on the number of proposals &
+ the number of detected objects.
+
+ Args:
+ model: a detectron2 model that takes `list[dict]` as input.
+ inputs (list[dict]): inputs to model, in detectron2's standard format.
+ """
+ return _wrapper_count_operators(model=model, inputs=inputs, mode=ACTIVATIONS_MODE, **kwargs)
+
+
+def _wrapper_count_operators(
+ model: nn.Module, inputs: list, mode: str, **kwargs
+) -> typing.DefaultDict[str, float]:
+
+ assert len(inputs) == 1, "Please use batch size=1"
tensor_input = inputs[0]["image"]
class WrapModel(nn.Module):
@@ -61,15 +97,21 @@ def flop_count_operators(
else:
log_first_n(
logging.WARN,
- f"Field '{k}' in output instances is not included in flops counting.",
+ f"Field '{k}' in output instances is not included"
+ " in flops/activations count.",
n=10,
)
return tuple(ret)
- raise NotImplementedError("Flops for segmentation models not supported yet.")
+ raise NotImplementedError("Count for segmentation models is not supported yet.")
old_train = model.training
with torch.no_grad():
+ if mode == FLOPS_MODE:
ret = flop_count(WrapModel(model).train(False), (tensor_input,), **kwargs)
+ elif mode == ACTIVATIONS_MODE:
+ ret = activation_count(WrapModel(model).train(False), (tensor_input,), **kwargs)
+ else:
+ raise NotImplementedError("Count for mode {} is not supported yet.".format(mode))
# compatible with change in fvcore
if isinstance(ret, tuple):
ret = ret[0]
