From 675e60d5a5c6120ca01ed77cb8a7300c150e6b76 Mon Sep 17 00:00:00 2001
From: gaotingquan <gaotingquan@baidu.com>
Date: Tue, 17 May 2022 08:44:30 +0000
Subject: [PATCH] feat: log the FLOPs and Params when building model

---
 ppcls/arch/__init__.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/ppcls/arch/__init__.py b/ppcls/arch/__init__.py
index da21e101..a718e11b 100644
--- a/ppcls/arch/__init__.py
+++ b/ppcls/arch/__init__.py
@@ -15,6 +15,7 @@
 import copy
 import importlib
 
+import paddle
 import paddle.nn as nn
 from paddle.jit import to_static
 from paddle.static import InputSpec
@@ -40,6 +41,15 @@ def build_model(config):
     if isinstance(arch, TheseusLayer):
         prune_model(config, arch)
         quantize_model(config, arch)
+
+    logger.info("The FLOPs and Params of Arch:")
+    try:
+        flops = paddle.flops(arch, [1, *config["Global"]["image_shape"]])
+    except Exception as e:
+        logger.warning(
+            f"An error occurred when calculating FLOPs and Params of Arch. Please check the Global.image_shape in config. The details of error is: {e}"
+        )
+
     return arch
 
 
-- 
GitLab