halib 0.1.91__py3-none-any.whl → 0.2.21__py3-none-any.whl

halib/research/perf/flop_calc.py

@@ -0,0 +1,190 @@
+import os
+import sys
+import torch
+import timm
+from argparse import ArgumentParser
+from fvcore.nn import FlopCountAnalysis
+from halib import *
+from halib.filetype import csvfile
+from curriculum.utils.config import *
+from curriculum.utils.model_helper import *
+
+
+# ---------------------------------------------------------------------
+# Argument Parser
+# ---------------------------------------------------------------------
+def parse_args():
+    parser = ArgumentParser(description="Calculate FLOPs for TIMM or trained models")
+
+    # Option 1: Direct TIMM model
+    parser.add_argument(
+        "--model_name", type=str, help="TIMM model name (e.g., efficientnet_b0)"
+    )
+    parser.add_argument(
+        "--num_classes", type=int, default=1000, help="Number of output classes"
+    )
+
+    # Option 2: Experiment directory
+    parser.add_argument(
+        "--indir",
+        type=str,
+        default=None,
+        help="Directory containing trained experiment (with .yaml and .pth)",
+    )
+    parser.add_argument(
+        "-o", "--o", action="store_true", help="Open output CSV after saving"
+    )
+    return parser.parse_args()
+
+
+# ---------------------------------------------------------------------
+# Helper Functions
+# ---------------------------------------------------------------------
+def _get_list_of_proc_dirs(indir):
+    assert os.path.exists(indir), f"Input directory {indir} does not exist."
+    pth_files = [f for f in os.listdir(indir) if f.endswith(".pth")]
+    if len(pth_files) > 0:
+        return [indir]
+    return [
+        os.path.join(indir, f)
+        for f in os.listdir(indir)
+        if os.path.isdir(os.path.join(indir, f))
+    ]
+
+
+def _calculate_flops_for_model(model_name, num_classes):
+    """Calculate FLOPs for a plain TIMM model."""
+    try:
+        model = timm.create_model(model_name, pretrained=False, num_classes=num_classes)
+        input_size = timm.data.resolve_data_config(model.default_cfg)["input_size"]
+        dummy_input = torch.randn(1, *input_size)
+        model.eval() # ! set to eval mode to avoid some warnings or errors
+        flops = FlopCountAnalysis(model, dummy_input)
+        gflops = flops.total() / 1e9
+        mflops = flops.total() / 1e6
+        print(f"\nModel: **{model_name}**, Classes: {num_classes}")
+        print(f"Input size: {input_size}, FLOPs: **{gflops:.3f} GFLOPs**, **{mflops:.3f} MFLOPs**\n")
+        return model_name, gflops, mflops
+    except Exception as e:
+        print(f"[Error] Could not calculate FLOPs for {model_name}: {e}")
+        return model_name, -1, -1
+
+
+def _calculate_flops_for_experiment(exp_dir):
+    """Calculate FLOPs for a trained experiment directory."""
+    yaml_files = [f for f in os.listdir(exp_dir) if f.endswith(".yaml")]
+    pth_files = [f for f in os.listdir(exp_dir) if f.endswith(".pth")]
+
+    assert (
+        len(yaml_files) == 1
+    ), f"Expected 1 YAML file in {exp_dir}, found {len(yaml_files)}"
+    assert (
+        len(pth_files) == 1
+    ), f"Expected 1 PTH file in {exp_dir}, found {len(pth_files)}"
+
+    exp_cfg_yaml = os.path.join(exp_dir, yaml_files[0])
+    cfg = ExpConfig.from_yaml(exp_cfg_yaml)
+    ds_label_list = cfg.dataset.get_label_list()
+
+    try:
+        model = build_model(
+            cfg.model.name, num_classes=len(ds_label_list), pretrained=True
+        )
+        model_weights_path = os.path.join(exp_dir, pth_files[0])
+        model.load_state_dict(torch.load(model_weights_path, map_location="cpu"))
+        model.eval()
+
+        input_size = timm.data.resolve_data_config(model.default_cfg)["input_size"]
+        dummy_input = torch.randn(1, *input_size)
+        flops = FlopCountAnalysis(model, dummy_input)
+        gflops = flops.total() / 1e9
+        mflops = flops.total() / 1e6
+
+        return str(cfg), cfg.model.name, gflops, mflops
+    except Exception as e:
+        console.print(f"[red] Error processing {exp_dir}: {e}[/red]")
+        return str(cfg), cfg.model.name, -1, -1
+
+
+# ---------------------------------------------------------------------
+# Main Entry
+# ---------------------------------------------------------------------
+def main():
+    args = parse_args()
+
+    # Case 1: Direct TIMM model input
+    if args.model_name:
+        _calculate_flops_for_model(args.model_name, args.num_classes)
+        return
+
+    # Case 2: Experiment directory input
+    if args.indir is None:
+        print("[Error] Either --model_name or --indir must be specified.")
+        return
+
+    proc_dirs = _get_list_of_proc_dirs(args.indir)
+    pprint(proc_dirs)
+
+    dfmk = csvfile.DFCreator()
+    TABLE_NAME = "model_flops_results"
+    dfmk.create_table(TABLE_NAME, ["exp_name", "model_name", "gflops", "mflops"])
+
+    console.rule(f"Calculating FLOPs for models in {len(proc_dirs)} dir(s)...")
+    rows = []
+    for exp_dir in tqdm(proc_dirs):
+        dir_name = os.path.basename(exp_dir)
+        console.rule(f"{dir_name}")
+        exp_name, model_name, gflops, mflops = _calculate_flops_for_experiment(exp_dir)
+        rows.append([exp_name, model_name, gflops, mflops])
+
+    dfmk.insert_rows(TABLE_NAME, rows)
+    dfmk.fill_table_from_row_pool(TABLE_NAME)
+
+    outfile = f"zout/zreport/{now_str()}_model_flops_results.csv"
+    dfmk[TABLE_NAME].to_csv(outfile, sep=";", index=False)
+    csvfile.fn_display_df(dfmk[TABLE_NAME])
+
+    if args.o:
+        os.system(f"start {outfile}")
+
+
+# ---------------------------------------------------------------------
+# Script Entry
+# ---------------------------------------------------------------------
+# flop_csv.py
+# if __name__ == "__main__":
+#     sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+#     main()
+
+
+# def main():
+#     csv_file = "./results-imagenet.csv"
+#     df = pd.read_csv(csv_file)
+#     # make param_count column as float
+#     # df['param_count'] = df['param_count'].astype(float)
+#     df["param_count"] = (
+#         pd.to_numeric(df["param_count"], errors="coerce").fillna(99999).astype(float)
+#     )
+#     df = df[df["param_count"] < 5.0]  # filter models with param_count < 20M
+
+#     dict_ls = []
+
+#     for index, row in tqdm(df.iterrows()):
+#         console.rule(f"Row {index+1}/{len(df)}")
+#         model = row["model"]
+#         num_class = 2
+#         _, _, mflops = _calculate_flops_for_model(model, num_class)
+#         dict_ls.append(
+#             {"model": model, "param_count": row["param_count"], "mflops": mflops}
+#         )
+
+#     # Create a DataFrame from the list of dictionaries
+#     result_df = pd.DataFrame(dict_ls)
+
+#     final_df = pd.merge(df, result_df, on=["model", "param_count"])
+#     final_df.sort_values(by="mflops", inplace=True, ascending=True)
+#     csvfile.fn_display_df(final_df)
+
+
+# if __name__ == "__main__":
+#     main()

halib/research/perf/gpu_mon.py

@@ -0,0 +1,58 @@
+# install `pynvml_utils` package first
+# see this repo: https://github.com/gpuopenanalytics/pynvml
+from pynvml_utils import nvidia_smi
+import time
+import threading
+from rich.pretty import pprint
+
+class GPUMonitor:
+    def __init__(self, gpu_index=0, interval=0.01):
+        self.nvsmi = nvidia_smi.getInstance()
+        self.gpu_index = gpu_index
+        self.interval = interval
+        self.gpu_stats = []
+        self._running = False
+        self._thread = None
+
+    def _monitor(self):
+        while self._running:
+            stats = self.nvsmi.DeviceQuery("power.draw, memory.used")["gpu"][
+                self.gpu_index
+            ]
+            # pprint(stats)
+            self.gpu_stats.append(
+                {
+                    "power": stats["power_readings"]["power_draw"],
+                    "power_unit": stats["power_readings"]["unit"],
+                    "memory": stats["fb_memory_usage"]["used"],
+                    "memory_unit": stats["fb_memory_usage"]["unit"],
+                }
+            )
+            time.sleep(self.interval)
+
+    def start(self):
+        if not self._running:
+            self._running = True
+            # clear previous stats
+            self.gpu_stats.clear()
+            self._thread = threading.Thread(target=self._monitor)
+            self._thread.start()
+
+    def stop(self):
+        if self._running:
+            self._running = False
+            self._thread.join()
+            # clear the thread reference
+            self._thread = None
+
+    def get_stats(self):
+        ## return self.gpu_stats
+        assert self._running is False, "GPU monitor is still running. Stop it first."
+
+        powers = [s["power"] for s in self.gpu_stats if s["power"] is not None]
+        memories = [s["memory"] for s in self.gpu_stats if s["memory"] is not None]
+        avg_power = sum(powers) / len(powers) if powers else 0
+        max_memory = max(memories) if memories else 0
+        # power_unit = self.gpu_stats[0]["power_unit"] if self.gpu_stats else "W"
+        # memory_unit = self.gpu_stats[0]["memory_unit"] if self.gpu_stats else "MiB"
+        return {"gpu_avg_power": avg_power, "gpu_avg_max_memory": max_memory}

halib/research/perf/perfcalc.py

@@ -0,0 +1,363 @@
+import os
+import glob
+from typing import Optional, Tuple
+import pandas as pd
+
+from abc import ABC, abstractmethod
+from collections import OrderedDict
+
+
+from ...common.common import now_str
+from ...system import filesys as fs
+
+from .perftb import PerfTB
+from .perfmetrics import *
+
+
+REQUIRED_COLS = ["experiment", "dataset"]
+CSV_FILE_POSTFIX = "__perf"
+METRIC_PREFIX = "metric_"
+
+
+class PerfCalc(ABC):  # Abstract base class for performance calculation
+    @abstractmethod
+    def get_experiment_name(self) -> str:
+        """
+        Return the name of the experiment.
+        This function should be overridden by the subclass if needed.
+        """
+        pass
+
+    @abstractmethod
+    def get_dataset_name(self) -> str:
+        """
+        Return the name of the dataset.
+        This function should be overridden by the subclass if needed.
+        """
+        pass
+
+    @abstractmethod
+    def get_metric_backend(self) -> MetricsBackend:
+        """
+        Return a list of metric names to be used for performance calculation OR a dictionaray with keys as metric names and values as metric instances of torchmetrics.Metric. For example: {"accuracy": Accuracy(), "precision": Precision()}
+
+        """
+        pass
+
+    def valid_proc_extra_data(self, proc_extra_data):
+        # make sure that all items in proc_extra_data are dictionaries, with same keys
+        if proc_extra_data is None or len(proc_extra_data) == 0:
+            return
+        if not all(isinstance(item, dict) for item in proc_extra_data):
+            raise TypeError("All items in proc_extra_data must be dictionaries")
+
+        if not all(
+            item.keys() == proc_extra_data[0].keys() for item in proc_extra_data
+        ):
+            raise ValueError(
+                "All dictionaries in proc_extra_data must have the same keys"
+            )
+
+    def valid_proc_metric_raw_data(self, metric_names, proc_metric_raw_data):
+        # make sure that all items in proc_metric_raw_data are dictionaries, with same keys as metric_names
+        assert (
+            isinstance(proc_metric_raw_data, list) and len(proc_metric_raw_data) > 0
+        ), "raw_data_for_metrics must be a non-empty list of dictionaries"
+
+        # make sure that all items in proc_metric_raw_data are dictionaries with keys as metric_names
+        if not all(isinstance(item, dict) for item in proc_metric_raw_data):
+            raise TypeError("All items in raw_data_for_metrics must be dictionaries")
+        if not all(
+            set(item.keys()) == set(metric_names) for item in proc_metric_raw_data
+        ):
+            raise ValueError(
+                "All dictionaries in raw_data_for_metrics must have the same keys as metric_names"
+            )
+
+    # ! only need to override this method if torchmetrics are not used
+    def calc_exp_perf_metrics(
+        self, metric_names, raw_metrics_data, extra_data=None, *args, **kwargs
+    ):
+        assert isinstance(raw_metrics_data, dict) or isinstance(
+            raw_metrics_data, list
+        ), "raw_data_for_metrics must be a dictionary or a list"
+
+        if extra_data is not None:
+            assert isinstance(
+                extra_data, type(raw_metrics_data)
+            ), "extra_data must be of the same type as raw_data_for_metrics (dict or list)"
+        # prepare raw_metric data for processing
+        proc_metric_raw_data_ls = (
+            raw_metrics_data
+            if isinstance(raw_metrics_data, list)
+            else [raw_metrics_data.copy()]
+        )
+        self.valid_proc_metric_raw_data(metric_names, proc_metric_raw_data_ls)
+        # prepare extra data for processing
+        proc_extra_data_ls = []
+        if extra_data is not None:
+            proc_extra_data_ls = (
+                extra_data if isinstance(extra_data, list) else [extra_data.copy()]
+            )
+            assert len(proc_extra_data_ls) == len(
+                proc_metric_raw_data_ls
+            ), "extra_data must have the same length as raw_data_for_metrics if it is a list"
+        # validate the extra_data
+        self.valid_proc_extra_data(proc_extra_data_ls)
+
+        # calculate the metrics output results
+        metrics_backend = self.get_metric_backend()
+        proc_outdict_list = []
+        for idx, raw_metrics_data in enumerate(proc_metric_raw_data_ls):
+            out_dict = {
+                "dataset": self.get_dataset_name(),
+                "experiment": self.get_experiment_name(),
+            }
+            custom_fields = []
+            if len(proc_extra_data_ls) > 0:
+                # add extra data to the output dictionary
+                extra_data_item = proc_extra_data_ls[idx]
+                out_dict.update(extra_data_item)
+                custom_fields = list(extra_data_item.keys())
+            metric_results = metrics_backend.calc_metrics(
+                metrics_data_dict=raw_metrics_data, *args, **kwargs
+            )
+            metric_results_prefix = {
+                f"metric_{k}": v for k, v in metric_results.items()
+            }
+            out_dict.update(metric_results_prefix)
+            ordered_cols = (
+                REQUIRED_COLS + custom_fields + list(metric_results_prefix.keys())
+            )
+            out_dict = OrderedDict(
+                (col, out_dict[col]) for col in ordered_cols if col in out_dict
+            )
+            proc_outdict_list.append(out_dict)
+
+        return proc_outdict_list
+
+    #! custom kwargs:
+    #! outfile - if provided, will save the output to a CSV file with the given path
+    #! outdir - if provided, will save the output to a CSV file in the given directory with a generated filename
+    #! return_df - if True, will return a DataFrame instead of a dictionary
+    def calc_perfs(
+        self,
+        raw_metrics_data: Union[List[dict], dict],
+        extra_data: Optional[Union[List[dict], dict]] = None,
+        *args,
+        **kwargs,
+    ) -> Tuple[Union[List[OrderedDict], pd.DataFrame], Optional[str]]:
+        """
+        Calculate the metrics.
+        This function should be overridden by the subclass if needed.
+        Must return a dictionary with keys as metric names and values as the calculated metrics.
+        """
+        metric_names = self.get_metric_backend().metric_names
+        out_dict_list = self.calc_exp_perf_metrics(
+            metric_names=metric_names,
+            raw_metrics_data=raw_metrics_data,
+            extra_data=extra_data,
+            *args,
+            **kwargs,
+        )
+        csv_outfile = kwargs.get("outfile", None)
+        if csv_outfile is not None:
+            filePathNoExt, _ = os.path.splitext(csv_outfile)
+            # pprint(f"CSV Outfile Path (No Ext): {filePathNoExt}")
+            csv_outfile = f"{filePathNoExt}{CSV_FILE_POSTFIX}.csv"
+        elif "outdir" in kwargs:
+            csvoutdir = kwargs["outdir"]
+            csvfilename = f"{now_str()}_{self.get_dataset_name()}_{self.get_experiment_name()}_{CSV_FILE_POSTFIX}.csv"
+            csv_outfile = os.path.join(csvoutdir, csvfilename)
+
+        # convert out_dict to a DataFrame
+        df = pd.DataFrame(out_dict_list)
+        # get the orders of the columns as the orders or the keys in out_dict
+        ordered_cols = list(out_dict_list[0].keys())
+        df = df[ordered_cols]  # reorder columns
+        if csv_outfile:
+            df.to_csv(csv_outfile, index=False, sep=";", encoding="utf-8")
+        return_df = kwargs.get("return_df", False)
+        if return_df:  # return DataFrame instead of dict if requested
+            return df, csv_outfile
+        else:
+            return out_dict_list, csv_outfile
+
+    @staticmethod
+    def default_exp_csv_filter_fn(exp_file_name: str) -> bool:
+        """
+        Default filter function for experiments.
+        Returns True if the experiment name does not start with "test_" or "debug_".
+        """
+        return "__perf.csv" in exp_file_name
+
+    @classmethod
+    def get_perftb_for_multi_exps(
+        cls,
+        indir: str,
+        exp_csv_filter_fn=default_exp_csv_filter_fn,
+        include_file_name=False,
+        csv_sep=";",
+    ) -> PerfTB:
+        """
+        Generate a performance report by scanning experiment subdirectories.
+        Must return a dictionary with keys as metric names and values as performance tables.
+        """
+
+        def get_df_for_all_exp_perf(csv_perf_files, csv_sep=";"):
+            """
+            Create a single DataFrame from all CSV files.
+            Assumes all CSV files MAY have different metrics
+            """
+            cols = []
+            FILE_NAME_COL = "file_name" if include_file_name else None
+
+            for csv_file in csv_perf_files:
+                temp_df = pd.read_csv(csv_file, sep=csv_sep)
+                if FILE_NAME_COL:
+                    temp_df[FILE_NAME_COL] = fs.get_file_name(
+                        csv_file, split_file_ext=False
+                    )
+                    # csvfile.fn_display_df(temp_df)
+                temp_df_cols = temp_df.columns.tolist()
+                for col in temp_df_cols:
+                    if col not in cols:
+                        cols.append(col)
+
+            df = pd.DataFrame(columns=cols)
+            for csv_file in csv_perf_files:
+                temp_df = pd.read_csv(csv_file, sep=csv_sep)
+                if FILE_NAME_COL:
+                    temp_df[FILE_NAME_COL] = fs.get_file_name(
+                        csv_file, split_file_ext=False
+                    )
+                # Drop all-NA columns to avoid dtype inconsistency
+                temp_df = temp_df.dropna(axis=1, how="all")
+                # ensure all columns are present in the final DataFrame
+                for col in cols:
+                    if col not in temp_df.columns:
+                        temp_df[col] = None  # fill missing columns with None
+                df = pd.concat([df, temp_df], ignore_index=True)
+            # assert that REQUIRED_COLS are present in the DataFrame
+            # pprint(df.columns.tolist())
+            sticky_cols = REQUIRED_COLS + (
+                [FILE_NAME_COL] if include_file_name else []
+            )  # columns that must always be present
+            for col in sticky_cols:
+                if col not in df.columns:
+                    raise ValueError(
+                        f"Required column '{col}' is missing from the DataFrame. REQUIRED_COLS = {sticky_cols}"
+                    )
+            metric_cols = [col for col in df.columns if col.startswith(METRIC_PREFIX)]
+            assert (
+                len(metric_cols) > 0
+            ), "No metric columns found in the DataFrame. Ensure that the CSV files contain metric columns starting with 'metric_'."
+            final_cols = sticky_cols + metric_cols
+            df = df[final_cols]
+            # # !hahv debug
+            # pprint("------ Final DataFrame Columns ------")
+            # csvfile.fn_display_df(df)
+            # ! validate all rows in df before returning
+            # make sure all rows will have at least values for REQUIRED_COLS and at least one metric column
+            for index, row in df.iterrows():
+                if not all(col in row and pd.notna(row[col]) for col in sticky_cols):
+                    raise ValueError(
+                        f"Row {index} is missing required columns or has NaN values in required columns: {row}"
+                    )
+                if not any(pd.notna(row[col]) for col in metric_cols):
+                    raise ValueError(f"Row {index} has no metric values: {row}")
+            # make sure these is no (experiment, dataset) pair that is duplicated
+            duplicates = df.duplicated(subset=sticky_cols, keep=False)
+            if duplicates.any():
+                raise ValueError(
+                    "Duplicate (experiment, dataset) pairs found in the DataFrame. Please ensure that each experiment-dataset combination is unique."
+                )
+            return df
+
+        def mk_perftb_report(df):
+            """
+            Create a performance report table from the DataFrame.
+            This function should be customized based on the specific requirements of the report.
+            """
+            perftb = PerfTB()
+            # find all "dataset" values (unique)
+            dataset_names = list(df["dataset"].unique())
+            # find all columns that start with METRIC_PREFIX
+            metric_cols = [col for col in df.columns if col.startswith(METRIC_PREFIX)]
+
+            # Determine which metrics are associated with each dataset.
+            # Since a dataset may appear in multiple rows and may not include all metrics in each, identify the row with the same dataset that contains the most non-NaN metric values. The set of metrics for that dataset is defined by the non-NaN metrics in that row.
+
+            dataset_metrics = {}
+            for dataset_name in dataset_names:
+                dataset_rows = df[df["dataset"] == dataset_name]
+                # Find the row with the most non-NaN metric values
+                max_non_nan_row = dataset_rows[metric_cols].count(axis=1).idxmax()
+                metrics_for_dataset = (
+                    dataset_rows.loc[max_non_nan_row, metric_cols]
+                    .dropna()
+                    .index.tolist()
+                )
+                dataset_metrics[dataset_name] = metrics_for_dataset
+
+            for dataset_name, metrics in dataset_metrics.items():
+                # Create a new row for the performance table
+                perftb.add_dataset(dataset_name, metrics)
+
+            for _, row in df.iterrows():
+                dataset_name = row["dataset"]
+                ds_metrics = dataset_metrics.get(dataset_name)
+                if dataset_name in dataset_metrics:
+                    # Add the metrics for this row to the performance table
+                    exp_name = row.get("experiment")
+                    exp_metric_values = {}
+                    for metric in ds_metrics:
+                        if metric in row and pd.notna(row[metric]):
+                            exp_metric_values[metric] = row[metric]
+                    perftb.add_experiment(
+                        experiment_name=exp_name,
+                        dataset_name=dataset_name,
+                        metrics=exp_metric_values,
+                    )
+
+            return perftb
+
+        assert os.path.exists(indir), f"Input directory {indir} does not exist."
+
+        csv_perf_files = []
+        # Find experiment subdirectories
+        exp_dirs = [
+            os.path.join(indir, d)
+            for d in os.listdir(indir)
+            if os.path.isdir(os.path.join(indir, d))
+        ]
+        if len(exp_dirs) == 0:
+            csv_perf_files = glob.glob(os.path.join(indir, f"*.csv"))
+            csv_perf_files = [
+                file_item
+                for file_item in csv_perf_files
+                if exp_csv_filter_fn(file_item)
+            ]
+        else:
+            # multiple experiment directories found
+            # Collect all matching CSV files in those subdirs
+            for exp_dir in exp_dirs:
+                # pprint(f"Searching in experiment directory: {exp_dir}")
+                matched = glob.glob(os.path.join(exp_dir, f"*.csv"))
+                matched = [
+                    file_item for file_item in matched if exp_csv_filter_fn(file_item)
+                ]
+                csv_perf_files.extend(matched)
+
+        assert (
+            len(csv_perf_files) > 0
+        ), f"No CSV files matching pattern '{exp_csv_filter_fn}' found in the experiment directories."
+
+        assert (
+            len(csv_perf_files) > 0
+        ), f"No CSV files matching pattern '{exp_csv_filter_fn}' found in the experiment directories."
+
+        all_exp_perf_df = get_df_for_all_exp_perf(csv_perf_files, csv_sep=csv_sep)
+        # csvfile.fn_display_df(all_exp_perf_df)
+        perf_tb = mk_perftb_report(all_exp_perf_df)
+        return perf_tb