From 4fa94d4ab4c2174688f1327259e51916d4428796 Mon Sep 17 00:00:00 2001
From: yukavio <67678385+yukavio@users.noreply.github.com>
Date: Fri, 25 Dec 2020 10:45:26 +0800
Subject: [PATCH] flops fix Cherry pick (#29872)

* add some feature for paddle.flops (#29572)

* fix flops (#29758)

* fix flops

* fix flops

* fix flops (#29818)
---
 python/paddle/hapi/dynamic_flops.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/python/paddle/hapi/dynamic_flops.py b/python/paddle/hapi/dynamic_flops.py
index 382227ea832..bfbb483ac31 100644
--- a/python/paddle/hapi/dynamic_flops.py
+++ b/python/paddle/hapi/dynamic_flops.py
@@ -221,7 +221,8 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False):
         if m_type in custom_ops:
             flops_fn = custom_ops[m_type]
             if m_type not in types_collection:
-                print("Customize Function has been appied to {}".format(m_type))
+                print("Customize Function has been applied to {}".format(
+                    m_type))
         elif m_type in register_hooks:
             flops_fn = register_hooks[m_type]
             if m_type not in types_collection:
@@ -229,7 +230,7 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False):
         else:
             if m_type not in types_collection:
                 print(
-                    "Cannot find suitable count function for {}. Treat it as zero Macs.".
+                    "Cannot find suitable count function for {}. Treat it as zero FLOPs.".
                     format(m_type))
 
         if flops_fn is not None:
@@ -254,11 +255,10 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False):
     for m in model.sublayers():
         if len(list(m.children())) > 0:
             continue
-        total_ops += m.total_ops
-        total_params += m.total_params
-
-    total_ops = int(total_ops)
-    total_params = int(total_params)
+        if set(['total_ops', 'total_params', 'input_shape',
+                'output_shape']).issubset(set(list(m._buffers.keys()))):
+            total_ops += m.total_ops
+            total_params += m.total_params
 
     if training:
         model.train()
@@ -277,7 +277,8 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False):
     for n, m in model.named_sublayers():
         if len(list(m.children())) > 0:
             continue
-        if "total_ops" in m._buffers:
+        if set(['total_ops', 'total_params', 'input_shape',
+                'output_shape']).issubset(set(list(m._buffers.keys()))):
             table.add_row([
                 m.full_name(), list(m.input_shape.numpy()),
                 list(m.output_shape.numpy()), int(m.total_params),
@@ -289,6 +290,6 @@ def dynamic_flops(model, inputs, custom_ops=None, print_detail=False):
             m._buffers.pop('output_shape')
     if (print_detail):
         print(table)
-    print('Total Flops: {}     Total Params: {}'.format(total_ops,
-                                                        total_params))
-    return total_ops
+    print('Total Flops: {}     Total Params: {}'.format(
+        int(total_ops), int(total_params)))
+    return int(total_ops)
-- 
GitLab