Add flops and params counter #2535
Conversation
|
Please remove accidentally pushed pth files. |
src/sdk/pynni/nni/counter.py
Outdated
| weight_mask = None | ||
| m_type = type(m) | ||
| if m_type in custom_ops: | ||
| if isinstance(m_list[idx-1], PrunerModuleWrapper): |
There was a problem hiding this comment.
What if idx == 0? What if PrunerModuleWrapper has multiple children?
There was a problem hiding this comment.
Thanks~ PrunerModuleWrapper is our customize wrapper and it directly wraps Conv or Linear module, so it won't have multiple children.
|
please add doc here https://nni.readthedocs.io/en/latest/Compressor/CompressionUtils.html, and put api reference here https://nni.readthedocs.io/en/latest/Compressor/CompressionReference.html |
src/sdk/pynni/nni/counter.py
Outdated
| if isinstance(prev_m, PrunerModuleWrapper): | ||
| weight_mask = prev_m.weight_mask | ||
|
|
||
| m.register_buffer('weight_mask', weight_mask) |
There was a problem hiding this comment.
no need to register buffer here if the model is not PrunerModuleWrapper?
There was a problem hiding this comment.
We do not directly register buffer on PrunerModuleWrapper. We only register buffer on Conv or Linear (please see custom_ops in Line 96) where the previous module was PrunerModuleWrapper.
| except ImportError: | ||
| _logger.warning('Please install thop using command: pip install thop') | ||
|
|
||
| def count_flops_params(model: nn.Module, input_size=None, verbose=True): |
There was a problem hiding this comment.
suggest to remove default value of input_size
There was a problem hiding this comment.
fix it. please review it on the latest version~
|
Please remove the pth files. |
Support flops and params calculation with mask.