halib 0.2.30__py3-none-any.whl

halib/exp/data/torchloader.py

@@ -0,0 +1,165 @@
+"""
+ * @author Hoang Van-Ha
+ * @email hoangvanhauit@gmail.com
+ * @create date 2024-03-27 15:40:22
+ * @modify date 2024-03-27 15:40:22
+ * @desc this module works as a utility tools for finding the best configuration for dataloader (num_workers, batch_size, pin_menory, etc.) that fits your hardware.
+"""
+from argparse import ArgumentParser
+
+import os
+import time
+import traceback
+
+from tqdm import tqdm
+from rich import inspect
+from typing import Union
+import itertools as it  # for cartesian product
+
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+
+from ...common.common import *
+from ...filetype import csvfile
+from ...filetype.yamlfile import load_yaml
+
+def parse_args():
+    parser = ArgumentParser(description="desc text")
+    parser.add_argument("-cfg", "--cfg", type=str, help="cfg file for searching")
+    return parser.parse_args()
+
+
+def get_test_range(cfg: dict, search_item="num_workers"):
+    item_search_cfg = cfg["search_space"].get(search_item, None)
+    if item_search_cfg is None:
+        raise ValueError(f"search_item: {search_item} not found in cfg")
+    if isinstance(item_search_cfg, list):
+        return item_search_cfg
+    elif isinstance(item_search_cfg, dict):
+        if "mode" in item_search_cfg:
+            mode = item_search_cfg["mode"]
+            assert mode in ["range", "list"], f"mode: {mode} not supported"
+            value_in_mode = item_search_cfg.get(mode, None)
+            if value_in_mode is None:
+                raise ValueError(f"mode<{mode}>: data not found in <{search_item}>")
+            if mode == "range":
+                assert len(value_in_mode) == 3, f"range must have 3 values: start, stop, step"
+                start = value_in_mode[0]
+                stop = value_in_mode[1]
+                step = value_in_mode[2]
+                return list(range(start, stop, step))
+            elif mode == "list":
+                return item_search_cfg["list"]
+    else:
+        return [item_search_cfg]  # for int, float, str, bool, etc.
+
+
+def load_an_batch(loader_iter):
+    start = time.time()
+    next(loader_iter)
+    end = time.time()
+    return end - start
+
+
+def test_dataloader_with_cfg(origin_dataloader: DataLoader, cfg: Union[dict, str]):
+    try:
+        if isinstance(cfg, str):
+            cfg = load_yaml(cfg, to_dict=True)
+        dfmk = csvfile.DFCreator()
+        search_items = ["batch_size", "num_workers", "persistent_workers", "pin_memory"]
+        batch_limit = cfg["general"]["batch_limit"]
+        csv_cfg = cfg["general"]["to_csv"]
+        log_batch_info = cfg["general"]["log_batch_info"]
+
+        save_to_csv = csv_cfg["enabled"]
+        log_dir = csv_cfg["log_dir"]
+        filename = csv_cfg["filename"]
+        filename = f"{now_str()}_{filename}.csv"
+        outfile = os.path.join(log_dir, filename)
+
+        dfmk.create_table(
+            "cfg_search",
+            (search_items + ["avg_time_taken"]),
+        )
+        ls_range_test = []
+        for item in search_items:
+            range_test = get_test_range(cfg, search_item=item)
+            range_test = [(item, i) for i in range_test]
+            ls_range_test.append(range_test)
+
+        all_combinations = list(it.product(*ls_range_test))
+
+        rows = []
+        for cfg_idx, combine in enumerate(all_combinations):
+            console.rule(f"Testing cfg {cfg_idx+1}/{len(all_combinations)}")
+            inspect(combine)
+            batch_size = combine[search_items.index("batch_size")][1]
+            num_workers = combine[search_items.index("num_workers")][1]
+            persistent_workers = combine[search_items.index("persistent_workers")][1]
+            pin_memory = combine[search_items.index("pin_memory")][1]
+
+            test_dataloader = DataLoader(origin_dataloader.dataset, batch_size=batch_size, num_workers=num_workers, persistent_workers=persistent_workers, pin_memory=pin_memory, shuffle=True)
+            row = [
+                batch_size,
+                num_workers,
+                persistent_workers,
+                pin_memory,
+                0.0,
+            ]
+
+            # calculate the avg time taken to load the data for <batch_limit> batches
+            trainiter = iter(test_dataloader)
+            time_elapsed = 0
+            pprint('Start testing...')
+            for i in tqdm(range(batch_limit)):
+                single_batch_time = load_an_batch(trainiter)
+                if log_batch_info:
+                    pprint(f"Batch {i+1} took {single_batch_time:.4f} seconds to load")
+                time_elapsed += single_batch_time
+            row[-1] = time_elapsed / batch_limit
+            rows.append(row)
+        dfmk.insert_rows('cfg_search', rows)
+        dfmk.fill_table_from_row_pool('cfg_search')
+        with ConsoleLog("results"):
+            csvfile.fn_display_df(dfmk['cfg_search'])
+            if save_to_csv:
+                dfmk["cfg_search"].to_csv(outfile, index=False)
+                console.print(f"[red] Data saved to <{outfile}> [/red]")
+
+    except Exception as e:
+        traceback.print_exc()
+        print(e)
+        # get current directory of this python file
+        current_dir = os.path.dirname(os.path.realpath(__file__))
+        standar_cfg_path = os.path.join(current_dir, "torchloader_search.yaml")
+        pprint(
+            f"Make sure you get the  right <cfg.yaml> file. An example of <cfg.yaml> file can be found at this path: {standar_cfg_path}"
+        )
+        return
+
+def main():
+    args = parse_args()
+    cfg_yaml = args.cfg
+    cfg_dict = load_yaml(cfg_yaml, to_dict=True)
+
+    # Define transforms for data augmentation and normalization
+    transform = transforms.Compose(
+        [
+            transforms.RandomHorizontalFlip(),  # Randomly flip images horizontally
+            transforms.RandomRotation(10),  # Randomly rotate images by 10 degrees
+            transforms.ToTensor(),  # Convert images to PyTorch tensors
+            transforms.Normalize(
+                (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
+            ),  # Normalize pixel values to [-1, 1]
+        ]
+    )
+    test_dataset = datasets.CIFAR10(
+        root="./data", train=False, download=True, transform=transform
+    )
+    batch_size = 64
+    train_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
+    test_dataloader_with_cfg(train_loader, cfg_dict)
+
+
+if __name__ == "__main__":
+    main()

halib/exp/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/exp/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}