Feature/extra apply funcs

.gitignore

@@ -0,0 +1 @@
+**/__pycache__/

calflops/flops_counter.py

@@ -30,6 +30,7 @@ def calculate_flops(model,
                     forward_mode="forward",
                     include_backPropagation=False,
                     compute_bp_factor=2.0,
+                    extra_apply_funcs=[],
                     print_results=True,
                     print_detailed=True,
                     output_as_string=True,
@@ -48,6 +49,7 @@ def calculate_flops(model,
         forward_mode (str, optional): To determine the mode of model inference, Default to 'forward'. And use 'generate' if model inference uses model.generate().
         include_backPropagation (bool, optional): Decides whether the final return FLOPs computation includes the computation for backpropagation.
         compute_bp_factor (float, optional): The model backpropagation is a multiple of the forward propagation computation. Default to 2.
+        extra_apply_funcs (bool, optional): A list of functions to use in model.apply() which will be run before print_results. Defaults to the empty list.
         print_results (bool, optional): Whether to print the model profile. Defaults to True.
         print_detailed (bool, optional): Whether to print the detailed model profile. Defaults to True.
         output_as_string (bool, optional): Whether to print the output as string. Defaults to True.
@@ -172,6 +174,10 @@ def calculate_flops(model,
     macs = calculate_flops_pipline.get_total_macs()
     params = calculate_flops_pipline.get_total_params()
 
+    if len(extra_apply_funcs) > 0:
+        for func in extra_apply_funcs:
+            model.apply(func)
+
     if print_results:
         return_print = calculate_flops_pipline.print_model_pipline(units=output_unit,
                                                                    precision=output_precision,

calflops/flops_counter_hf.py

@@ -33,6 +33,7 @@ def calculate_flops_hf(model_name,
                        forward_mode="forward",
                        include_backPropagation=False,
                        compute_bp_factor=2.0,
+                       extra_apply_funcs=[],
                        print_results=True,
                        print_detailed=True,
                        output_as_string=True,
@@ -51,6 +52,7 @@ def calculate_flops_hf(model_name,
         forward_mode (str, optional): To determine the mode of model inference, Default to 'forward'. And use 'generate' if model inference uses model.generate().
         include_backPropagation (bool, optional): Decides whether the final return FLOPs computation includes the computation for backpropagation.
         compute_bp_factor (float, optional): The model backpropagation is a multiple of the forward propagation computation. Default to 2.
+        extra_apply_funcs (bool, optional): A list of functions to use in model.apply() which will be run before print_results. Defaults to the empty list.
         print_results (bool, optional): Whether to print the model profile. Defaults to True.
         print_detailed (bool, optional): Whether to print the detailed model profile. Defaults to True.
         output_as_string (bool, optional): Whether to print the output as string. Defaults to True.
@@ -126,6 +128,10 @@ def calculate_flops_hf(model_name,
         params = calculate_flops_pipline.get_total_params()
 
 
+        if len(extra_apply_funcs) > 0:
+            for func in extra_apply_funcs:
+                model.apply(func)
+
         print_return = calculate_flops_pipline.print_return_model_pipline(units=output_unit,
                                                     precision=output_precision,
                                                     print_detailed=print_detailed,

pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "calflops"
+version = "0.3.2"
+authors = [
+  { name="MrYXJ", email="[email protected]"},
+]
+description = "This package is designed to compute the theoretical amount of FLOPs(floating-point operations)、MACs(multiply-add operations) and Parameters in all various neural networks, such as Linear、 CNN、 RNN、 GCN、Transformer(Bert、LlaMA etc Large Language Model)，including any custom models via torch.nn.function.* as long as based on the Pytorch implementation."
+
+readme = "README.md"
+license = { file="LICENSE" }
+requires-python = ">=3.6"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "torch>=1.1.0",
+    "accelerate>=0.22.0",
+    "huggingface_hub>=0.16.4"
+]
+
+[project.urls]
+"Homepage" = "https://github.com/MrYxJ/calculate-flops.pytorch"