Number of flops in a graph model?

[Chris-Pedersen]

Hi there, just wondering if anyone knows of a package/script to estimate the number of flops for a given graph model? It seems many of the standard packages for doing this such as fvcore.nn and pthflops do not work with torch.geometric models and datasets.

[rusty1s]

I am not an expert on this one, but can you share some information why these packages do not work with PyG?

[kitbransby]

The error message that I get from fvcore is:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[18], line 2
      1 flops = FlopCountAnalysis(model, image)
----> 2 flops.total()

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/fvcore/nn/jit_analysis.py:248, in JitModelAnalysis.total(self, module_name)
    237 def total(self, module_name: str = "") -> int:
    238     """
    239     Returns the total aggregated statistic across all operators
    240     for the requested module.
   (...)
    246         int : The aggregated statistic.
    247     """
--> 248     stats = self._analyze()
    249     module_name = self.canonical_module_name(module_name)
    250     total_count = sum(stats.counts[module_name].values())

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/fvcore/nn/jit_analysis.py:551, in JitModelAnalysis._analyze(self)
    549     elif self._warn_trace == "no_tracer_warning":
    550         warnings.filterwarnings("ignore", category=TracerWarning)
--> 551     graph = _get_scoped_trace_graph(self._model, self._inputs, self._aliases)
    553 # Assures even modules not in the trace graph are initialized to zero count
    554 counts = {}

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/fvcore/nn/jit_analysis.py:176, in _get_scoped_trace_graph(module, inputs, aliases)
    173     name = aliases[mod]
    174     register_hooks(mod, name)
--> 176 graph, _ = _get_trace_graph(module, inputs)
    178 for handle in hook_handles:
    179     handle.remove()

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/jit/_trace.py:1184, in _get_trace_graph(f, args, kwargs, strict, _force_outplace, return_inputs, _return_inputs_states)
   1182 if not isinstance(args, tuple):
   1183     args = (args,)
-> 1184 outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
   1185 return outs

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/jit/_trace.py:127, in ONNXTracedModule.forward(self, *args)
    124     else:
    125         return tuple(out_vars)
--> 127 graph, out = torch._C._create_graph_by_tracing(
    128     wrapper,
    129     in_vars + module_state,
    130     _create_interpreter_name_lookup_fn(),
    131     self.strict,
    132     self._force_outplace,
    133 )
    135 if self._return_inputs:
    136     return graph, outs[0], ret_inputs[0]

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/jit/_trace.py:118, in ONNXTracedModule.forward.<locals>.wrapper(*args)
    116 if self._return_inputs_states:
    117     inputs_states.append(_unflatten(in_args, in_desc))
--> 118 outs.append(self.inner(*trace_inputs))
    119 if self._return_inputs_states:
    120     inputs_states[0] = (inputs_states[0], trace_inputs)

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1212, in Module._call_impl(self, *input, **kwargs)
   1209     bw_hook = hooks.BackwardHook(self, full_backward_hooks)
   1210     input = bw_hook.setup_input_hook(input)
-> 1212 result = forward_call(*input, **kwargs)
   1213 if _global_forward_hooks or self._forward_hooks:
   1214     for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File XXXXXXXX, in Graph_IVUS.forward(self, x)
     95 ref_scaled = self.scale_points(ref_points, bbox_pred)
     97 #print(bbox_pred)
     98 #print(ref_scaled)
     99 
   (...)
    103 # end = torch.cuda.Event(enable_timing=True)
    104 # start.record()
--> 105 pts_pred = self.POLYGON_HEAD(feature_maps, ref_scaled)
    106 # end.record()
    107 # torch.cuda.synchronize()
    108 # print('polygon pred', start.elapsed_time(end))
    109 
    110 #print('gcn out: ', [o.shape for o in pts_pred])
    112 pts_pred = [o.reshape(bs, 2, -1, 2) for o in pts_pred]

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1212, in Module._call_impl(self, *input, **kwargs)
   1209     bw_hook = hooks.BackwardHook(self, full_backward_hooks)
   1210     input = bw_hook.setup_input_hook(input)
-> 1212 result = forward_call(*input, **kwargs)
   1213 if _global_forward_hooks or self._forward_hooks:
   1214     for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File /XXXXX, in GCN.forward(self, feat_maps, points)
    348 feat1 = self.sample_features(feat_maps, points)
    349 inp1 = torch.cat([feat1, points], -1)
--> 350 x1, mag1 = self.block1(inp1.reshape(bs, 16, -1))
    351 out1 = points.reshape(bs, -1, 2) + (self.step_size * norm1 * mag1)
    353 # BLOCK 2

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1212, in Module._call_impl(self, *input, **kwargs)
   1209     bw_hook = hooks.BackwardHook(self, full_backward_hooks)
   1210     input = bw_hook.setup_input_hook(input)
-> 1212 result = forward_call(*input, **kwargs)
   1213 if _global_forward_hooks or self._forward_hooks:
   1214     for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File XXXXX, in Graph_Block.forward(self, x)
    514 def forward(self, x):
    515     inp = x
--> 517     x = self.gcn1(x, self.edge_index)
    518     x = self.norm1(x)
    519     x = self.relu(x)

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1212, in Module._call_impl(self, *input, **kwargs)
   1209     bw_hook = hooks.BackwardHook(self, full_backward_hooks)
   1210     input = bw_hook.setup_input_hook(input)
-> 1212 result = forward_call(*input, **kwargs)
   1213 if _global_forward_hooks or self._forward_hooks:
   1214     for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch_geometric/nn/conv/gcn_conv.py:229, in GCNConv.forward(self, x, edge_index, edge_weight)
    226         else:
    227             edge_index = cache
--> 229 x = self.lin(x)
    231 # propagate_type: (x: Tensor, edge_weight: OptTensor)
    232 out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
    233                      size=None)

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1212, in Module._call_impl(self, *input, **kwargs)
   1209     bw_hook = hooks.BackwardHook(self, full_backward_hooks)
   1210     input = bw_hook.setup_input_hook(input)
-> 1212 result = forward_call(*input, **kwargs)
   1213 if _global_forward_hooks or self._forward_hooks:
   1214     for hook in (*_global_forward_hooks.values(), *self._forward_hooks.values()):

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch/nn/modules/module.py:1182, in Module._slow_forward(self, *input, **kwargs)
   1180         recording_scopes = False
   1181 try:
-> 1182     result = self.forward(*input, **kwargs)
   1183 finally:
   1184     if recording_scopes:

File ~/anaconda3/envs/torch_1.13_detr/lib/python3.10/site-packages/torch_geometric/nn/dense/linear.py:132, in Linear.forward(self, x)
    127 def forward(self, x: Tensor) -> Tensor:
    128     r"""
    129     Args:
    130         x (torch.Tensor): The input features.
    131     """
--> 132     return F.linear(x, self.weight, self.bias)

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
Tensor:

[rusty1s]

I can try to take a look soon :)

[rusty1s]

I fixed this in #7693, although I am not totally convinced fvcore computes FLOPS correctly, e.g., it seems to only be able to register matrix multiplications.