facebookresearch / facebookresearch/fvcore
Counting FLOPS for a custom op with set_op_handle: a toy example that doesn't work.
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 236
- PR merge metrics
- No merged PRs in 30d
Description
I'm extending the given example for [fvcore.nn.FlopCountAnalysis](https://detectron2.readthedocs.io/en/latest/modules/fvcore.html#fvcore.nn.FlopCountAnalysis) to add flops count of a custom op within my model class.
```
import torch
from collections import Counter
from fvcore.nn import FlopCountAnalysis
from torch import nn
class TestModel(nn.Module):
"""Toy model."""
def __init__(self):
super().__init__()
self.act = nn.ReLU()
self.conv = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=1)
self.fc = nn.Linear(in_features=1000, out_features=10)
def forward(self, x):
_ = self.custom_op_flop_counter(inputs=x, outputs=None)
return self.fc(self.act(self.conv(x)).flatten(1))
@staticmethod
# Has no access to anything else in the class.
def custom_op_flop_counter(inputs, outputs) -> Counter:
"""Returns counter value to include in flops."""
# The function should return a counter object with per-operator statistics.
return Counter({'custom_op': 500})
model = TestModel()
inputs = (torch.randn((1, 3, 10, 10)),)
flops = FlopCountAnalysis(
model,
inputs).set_op_handle(
"custom_op", model.custom_op_flop_counter)
print(flops.by_module_and_operator())
```
The "custom_op" and its returned value of `500` are not seen in the print statement. It does print the expected values for the `linear` and `conv` operators, e.g.: `{'': Counter({'linear': 10000, 'conv': 3000}), 'act': Counter(), 'conv': Counter({'conv': 3000}), 'fc': Counter({'linear': 10000})}`. What am I doing wrong here that prevents my custom op count being included?
Contributor guide
Assessment
This issue has not been assessed yet.