diff --git a/ppcls/arch/__init__.py b/ppcls/arch/__init__.py index da21e101a27eb0db2c05b658346148bda3139c80..a718e11b903430c386269f8b4add1298b51280b6 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