I'm extending the given example for 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?
I'm extending the given example for fvcore.nn.FlopCountAnalysis to add flops count of a custom op within my model class.
The "custom_op" and its returned value of
500
are not seen in the print statement. It does print the expected values for thelinear
andconv
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?