Merge pull request #120 from sovrasov/fix_linear_args

Fix parsing of args in interpolate hook

ptflops/pytorch_ops.py

@@ -312,7 +312,10 @@ def _numel_functional_flops_hook(input, *args, **kwargs):
 
 
 def _interpolate_functional_flops_hook(*args, **kwargs):
-    input = args[0]
+    input = kwargs.get('input', None)
+    if input is None and len(args) > 0:
+        input = args[0]
+
     size = kwargs.get('size', None)
     if size is None and len(args) > 1:
         size = args[1]

tests/common_test.py

@@ -60,3 +60,29 @@ def input_constructor(input_res):
                                       print_per_layer_stat=False)
 
         assert (macs, params) == (8, 8)
+
+    def test_func_interpolate_args(self):
+        class CustomModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return nn.functional.interpolate(input=x, size=(20, 20),
+                                                 mode='bilinear', align_corners=False)
+
+        macs, params = \
+            get_model_complexity_info(CustomModel(), (3, 10, 10),
+                                      as_strings=False,
+                                      print_per_layer_stat=False)
+        assert params == 0
+        assert macs > 0
+
+        CustomModel.forward = lambda self, x: nn.functional.interpolate(x, size=(20, 20),
+                                                                        mode='bilinear')
+
+        macs, params = \
+            get_model_complexity_info(CustomModel(), (3, 10, 10),
+                                      as_strings=False,
+                                      print_per_layer_stat=False)
+        assert params == 0
+        assert macs > 0