diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..61f2dc9 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +**/__pycache__/ diff --git a/calflops/__pycache__/__init__.cpython-311.pyc b/calflops/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index 5f37917..0000000 Binary files a/calflops/__pycache__/__init__.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/__init__.cpython-39.pyc b/calflops/__pycache__/__init__.cpython-39.pyc deleted file mode 100644 index f467f28..0000000 Binary files a/calflops/__pycache__/__init__.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/big_modeling.cpython-39.pyc b/calflops/__pycache__/big_modeling.cpython-39.pyc deleted file mode 100644 index b33536f..0000000 Binary files a/calflops/__pycache__/big_modeling.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/calculate_pipline.cpython-311.pyc b/calflops/__pycache__/calculate_pipline.cpython-311.pyc deleted file mode 100644 index 18b1b0c..0000000 Binary files a/calflops/__pycache__/calculate_pipline.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/calculate_pipline.cpython-39.pyc b/calflops/__pycache__/calculate_pipline.cpython-39.pyc deleted file mode 100644 index 588a6db..0000000 Binary files a/calflops/__pycache__/calculate_pipline.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/estimate.cpython-311.pyc b/calflops/__pycache__/estimate.cpython-311.pyc deleted file mode 100644 index dd59e89..0000000 Binary files a/calflops/__pycache__/estimate.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/estimate.cpython-39.pyc b/calflops/__pycache__/estimate.cpython-39.pyc deleted file mode 100644 index 186845c..0000000 Binary files a/calflops/__pycache__/estimate.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/estimate.py b/calflops/__pycache__/estimate.py deleted file mode 100644 index df179ec..0000000 --- a/calflops/__pycache__/estimate.py +++ /dev/null @@ -1,271 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2023 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -#import argparse - -from huggingface_hub import model_info -from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError - -from accelerate import init_empty_weights -from accelerate.utils import ( - # calculate_maximum_sizes, - # convert_bytes, - is_timm_available, - is_transformers_available, -) - - -if is_transformers_available(): - import transformers - from transformers import AutoConfig, AutoModel - -if is_timm_available(): - import timm - - -def verify_on_hub(repo: str, token: str = None): - "Verifies that the model is on the hub and returns the model info." - try: - return model_info(repo, token=token) - except GatedRepoError: - return "gated" - except RepositoryNotFoundError: - return "repo" - - -def check_has_model(error): - """ - Checks what library spawned `error` when a model is not found - """ - if is_timm_available() and isinstance(error, RuntimeError) and "Unknown model" in error.args[0]: - return "timm" - elif ( - is_transformers_available() - and isinstance(error, OSError) - and "does not appear to have a file named" in error.args[0] - ): - return "transformers" - else: - return "unknown" - - -def create_empty_model(model_name: str, library_name: str, trust_remote_code: bool = False, access_token: str = None): - """ - Creates an empty model from its parent library on the `Hub` to calculate the overall memory consumption. - - Args: - model_name (`str`): - The model name on the Hub - library_name (`str`): - The library the model has an integration with, such as `transformers`. Will be used if `model_name` has no - metadata on the Hub to determine the library. - trust_remote_code (`bool`, `optional`, defaults to `False`): - Whether or not to allow for custom models defined on the Hub in their own modeling files. This option - should only be set to `True` for repositories you trust and in which you have read the code, as it will - execute code present on the Hub on your local machine. - access_token (`str`, `optional`, defaults to `None`): - The access token to use to access private or gated models on the Hub. (for use on the Gradio app) - - Returns: - `torch.nn.Module`: The torch model that has been initialized on the `meta` device. - - """ - model_info = verify_on_hub(model_name, access_token) - # Simplified errors - if model_info == "gated": - raise GatedRepoError( - f"Repo for model `{model_name}` is gated. You must be authenticated to access it. Please run `huggingface-cli login`." - ) - elif model_info == "repo": - raise RepositoryNotFoundError( - f"Repo for model `{model_name}` does not exist on the Hub. If you are trying to access a private repo," - " make sure you are authenticated via `huggingface-cli login` and have access." - ) - if library_name is None: - library_name = getattr(model_info, "library_name", False) - if not library_name: - raise ValueError( - f"Model `{model_name}` does not have any library metadata on the Hub, please manually pass in a `--library_name` to use (such as `transformers`)" - ) - if library_name == "transformers": - if not is_transformers_available(): - raise ImportError( - f"To check `{model_name}`, `transformers` must be installed. Please install it via `pip install transformers`" - ) - print(f"Loading pretrained config for `{model_name}` from `transformers`...") - - #auto_map = model_info.config.get("auto_map", False) - auto_map = model_info.config.get("auto_map", True) - config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code) - - with init_empty_weights(): - # remote code could specify a specific `AutoModel` class in the `auto_map` - constructor = AutoModel - if isinstance(auto_map, dict): - value = None - for key in auto_map.keys(): - if key.startswith("AutoModelFor"): - value = key - break - if value is not None: - constructor = getattr(transformers, value) - model = constructor.from_config(config, trust_remote_code=trust_remote_code) - elif library_name == "timm": - if not is_timm_available(): - raise ImportError( - f"To check `{model_name}`, `timm` must be installed. Please install it via `pip install timm`" - ) - print(f"Loading pretrained config for `{model_name}` from `timm`...") - with init_empty_weights(): - model = timm.create_model(model_name, pretrained=False) - else: - raise ValueError( - f"Library `{library_name}` is not supported yet, please open an issue on GitHub for us to add support." - ) - return model - - -def create_ascii_table(headers: list, rows: list, title: str): - "Creates a pretty table from a list of rows, minimal version of `tabulate`." - sep_char, in_between = "│", "─" - column_widths = [] - for i in range(len(headers)): - column_values = [row[i] for row in rows] + [headers[i]] - max_column_width = max(len(value) for value in column_values) - column_widths.append(max_column_width) - - formats = [f"%{column_widths[i]}s" for i in range(len(rows[0]))] - - pattern = f"{sep_char}{sep_char.join(formats)}{sep_char}" - diff = 0 - - def make_row(left_char, middle_char, right_char): - return f"{left_char}{middle_char.join([in_between * n for n in column_widths])}{in_between * diff}{right_char}" - - separator = make_row("├", "┼", "┤") - if len(title) > sum(column_widths): - diff = abs(len(title) - len(separator)) - column_widths[-1] += diff - - # Update with diff - separator = make_row("├", "┼", "┤") - initial_rows = [ - make_row("┌", in_between, "┐"), - f"{sep_char}{title.center(len(separator) - 2)}{sep_char}", - make_row("├", "┬", "┤"), - ] - table = "\n".join(initial_rows) + "\n" - column_widths[-1] += diff - centered_line = [text.center(column_widths[i]) for i, text in enumerate(headers)] - table += f"{pattern % tuple(centered_line)}\n{separator}\n" - for i, line in enumerate(rows): - centered_line = [t.center(column_widths[i]) for i, t in enumerate(line)] - table += f"{pattern % tuple(centered_line)}\n" - table += f'└{"┴".join([in_between * n for n in column_widths])}┘' - - return table - - -# def estimate_command_parser(subparsers=None): -# if subparsers is not None: -# parser = subparsers.add_parser("estimate-memory") -# else: -# parser = argparse.ArgumentParser(description="Model size estimator for fitting a model onto CUDA memory.") - -# parser.add_argument("model_name", type=str, help="The model name on the Hugging Face Hub.") -# parser.add_argument( -# "--library_name", -# type=str, -# help="The library the model has an integration with, such as `transformers`, needed only if this information is not stored on the Hub.", -# choices=["timm", "transformers"], -# ) -# parser.add_argument( -# "--dtypes", -# type=str, -# nargs="+", -# default=["float32", "float16", "int8", "int4"], -# help="The dtypes to use for the model, must be one (or many) of `float32`, `float16`, `int8`, and `int4`", -# choices=["float32", "float16", "int8", "int4"], -# ) -# parser.add_argument( -# "--trust_remote_code", -# action="store_true", -# help="""Whether or not to allow for custom models defined on the Hub in their own modeling files. This flag -# should only be used for repositories you trust and in which you have read the code, as it will execute -# code present on the Hub on your local machine.""", -# ) - -# if subparsers is not None: -# parser.set_defaults(func=estimate_command) -# return parser - - -# def gather_data(args): -# "Creates an empty model and gathers the data for the sizes" -# try: -# model = create_empty_model( -# args.model_name, library_name=args.library_name, trust_remote_code=args.trust_remote_code -# ) -# except (RuntimeError, OSError) as e: -# library = check_has_model(e) -# if library != "unknown": -# raise RuntimeError( -# f"Tried to load `{args.model_name}` with `{library}` but a possible model to load was not found inside the repo." -# ) -# raise e - -# total_size, largest_layer = calculate_maximum_sizes(model) - -# data = [] - -# for dtype in args.dtypes: -# dtype_total_size = total_size -# dtype_largest_layer = largest_layer[0] -# if dtype == "float16": -# dtype_total_size /= 2 -# dtype_largest_layer /= 2 -# elif dtype == "int8": -# dtype_total_size /= 4 -# dtype_largest_layer /= 4 -# elif dtype == "int4": -# dtype_total_size /= 8 -# dtype_largest_layer /= 8 -# dtype_training_size = dtype_total_size * 4 -# data.append([dtype, dtype_largest_layer, dtype_total_size, dtype_training_size]) -# return data - - -# def estimate_command(args): -# data = gather_data(args) -# for row in data: -# for i, item in enumerate(row): -# if isinstance(item, (int, float)): -# row[i] = convert_bytes(item) - -# headers = ["dtype", "Largest Layer", "Total Size", "Training using Adam"] - -# title = f"Memory Usage for loading `{args.model_name}`" -# table = create_ascii_table(headers, data, title) -# print(table) - - -# def main(): -# parser = estimate_command_parser() -# args = parser.parse_args() -# estimate_command(args) - - -# if __name__ == "__main__": -# main() \ No newline at end of file diff --git a/calflops/__pycache__/flops_counter.cpython-311.pyc b/calflops/__pycache__/flops_counter.cpython-311.pyc deleted file mode 100644 index 08ce4fd..0000000 Binary files a/calflops/__pycache__/flops_counter.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/flops_counter.cpython-39.pyc b/calflops/__pycache__/flops_counter.cpython-39.pyc deleted file mode 100644 index 12aa21a..0000000 Binary files a/calflops/__pycache__/flops_counter.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/flops_counter_hf.cpython-311.pyc b/calflops/__pycache__/flops_counter_hf.cpython-311.pyc deleted file mode 100644 index 00f2939..0000000 Binary files a/calflops/__pycache__/flops_counter_hf.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/flops_counter_hf.cpython-39.pyc b/calflops/__pycache__/flops_counter_hf.cpython-39.pyc deleted file mode 100644 index 23f834e..0000000 Binary files a/calflops/__pycache__/flops_counter_hf.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/pytorch_ops.cpython-311.pyc b/calflops/__pycache__/pytorch_ops.cpython-311.pyc deleted file mode 100644 index 82e8992..0000000 Binary files a/calflops/__pycache__/pytorch_ops.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/pytorch_ops.cpython-39.pyc b/calflops/__pycache__/pytorch_ops.cpython-39.pyc deleted file mode 100644 index f7e5e79..0000000 Binary files a/calflops/__pycache__/pytorch_ops.cpython-39.pyc and /dev/null differ diff --git a/calflops/__pycache__/utils.cpython-311.pyc b/calflops/__pycache__/utils.cpython-311.pyc deleted file mode 100644 index 1fbd18c..0000000 Binary files a/calflops/__pycache__/utils.cpython-311.pyc and /dev/null differ diff --git a/calflops/__pycache__/utils.cpython-39.pyc b/calflops/__pycache__/utils.cpython-39.pyc deleted file mode 100644 index 52f08b4..0000000 Binary files a/calflops/__pycache__/utils.cpython-39.pyc and /dev/null differ diff --git a/calflops/flops_counter.py b/calflops/flops_counter.py index ba216f3..4a2c6a2 100644 --- a/calflops/flops_counter.py +++ b/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, diff --git a/calflops/flops_counter_hf.py b/calflops/flops_counter_hf.py index 5d4e21e..edfc42e 100644 --- a/calflops/flops_counter_hf.py +++ b/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, diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..571ebc2 --- /dev/null +++ b/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="yxj2017@gmail.com"}, +] +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" \ No newline at end of file diff --git a/test_examples/test_extra_func.py b/test_examples/test_extra_func.py new file mode 100644 index 0000000..6d383c2 --- /dev/null +++ b/test_examples/test_extra_func.py @@ -0,0 +1,47 @@ +# !usr/bin/env python +# -*- coding:utf-8 -*- + +''' + Description : + Version : 1.0 + Author : MrYXJ + Mail : yxj2017@gmail.com + Github : https://github.com/MrYxJ + Date : 2023-08-24 11:48:59 + LastEditTime : 2023-08-24 19:42:16 + Copyright (C) 2023 mryxj. All rights reserved. +''' + +from calflops import calculate_flops +from calflops.utils import get_module_flops +from transformers import AutoModel +from transformers import AutoTokenizer +from collections import deque +from functools import partial + +def store_flops(module, datastore=None): + data = dict() + children = list(module.named_children()) + if len(children) > 0: + datastore.rotate(len(children)) + for name, child in children: + typename, flops, child_data = datastore.popleft() + data[name] = (typename, flops, child_data,) + datastore.append((type(module).__name__, get_module_flops(module), data,)) + +myqueue = deque() + +batch_size = 1 +max_seq_length = 128 +model_name = "bert-base-uncased" + +model = AutoModel.from_pretrained(model_name) +tokenizer = AutoTokenizer.from_pretrained(model_name) + +flops, macs, params = calculate_flops(model=model, + extra_apply_funcs=[partial(store_flops, datastore=myqueue)], + print_detailed=False, + input_shape=(batch_size,max_seq_length), + transformer_tokenizer=tokenizer) +print("Bert(hfl/chinese-roberta-wwm-ext) FLOPs:%s MACs:%s Params:%s \n" %(flops, macs, params)) +print(myqueue)