search

facebookresearch / fvcore

Collection of common code that's shared among different research projects in FAIR computer vision team.
Apache License 2.0
2k stars 226 forks source link

Per module jit analysis for flop count. #53

Closed ericmintun closed 3 years ago

ericmintun commented 3 years ago

Here's an implementation of how to obtain the per-module information from the jit trace. The class in fvcore/nn/jit_analysis.py can return per-module and per-operator information for generic calculation handles acting on operators in the model's trace graph.

The trace calculation and analysis step has been moved from fvcore/nn/jit_handles.py to fvcore/nn/jit_analysis.py. Functions to construct specific counters for flops and activations are defined in fvcore/nn/flop_count.py and fvcore/nn/activation_count.py.

facebook-github-bot commented 3 years ago

@ericmintun merged this pull request in facebookresearch/fvcore@cf76bbf85437f85443d6ef0c97e920d1a069b5ce.