mmengine.analysis.flop_count¶
- mmengine.analysis.flop_count(model, inputs, supported_ops=None)[source]¶
Given a model and an input to the model, compute the per-operator Gflops of the given model.
Adopted from https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/flop_count.py
- Parameters
model (nn.Module) – The model to compute flop counts.
inputs (tuple) – Inputs that are passed to model to count flops. Inputs need to be in a tuple.
supported_ops (dict(str,Callable) or None) – provide additional handlers for extra ops, or overwrite the existing handlers for convolution and matmul and einsum. The key is operator name and the value is a function that takes (inputs, outputs) of the op. We count one Multiply-Add as one FLOP.
- Returns
A dictionary that records the number of gflops for each operation and a Counter that records the number of unsupported operations.
- Return type
tuple[defaultdict, Counter]