nntool 1.3.0__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

nntool/train/trainer.py

@@ -0,0 +1,92 @@
+import wandb
+import torch
+
+from typing import Dict, Set
+from pathlib import Path
+from abc import abstractmethod
+
+
+class BaseTrainer(object):
+    def __init__(self, output_folder: str, has_wandb_writer: bool = False) -> None:
+        self.output_folder = Path(output_folder)
+        self.output_folder.mkdir(exist_ok=True)
+
+        self.has_writer = has_wandb_writer
+        self._wandb_defined_metrics: Set[str] = set()
+        if self.has_writer:
+            self._init_wandb_metric()
+
+    def _get_wandb_step_name(self, wandb_section: str = "trainer_state") -> str:
+        return f"{wandb_section}/{self.__class__.__name__}_step"
+
+    def _init_wandb_metric(self):
+        wandb.define_metric(self._get_wandb_step_name())
+        self._wandb_defined_metrics.add(self._get_wandb_step_name())
+
+    def log(self, log_dict: dict, step: int = None, section: str = "train"):
+        if not self.has_writer:
+            return
+
+        # add section to each metric
+        log_dict = {f"{section}/{k}": v for k, v in log_dict.items()}
+
+        # define metrics against custom step name
+        for k in log_dict.keys():
+            if k not in self._wandb_defined_metrics:
+                wandb.define_metric(k, step_metric=self._get_wandb_step_name())
+                self._wandb_defined_metrics.add(k)
+
+        # add step metrics
+        if step is not None:
+            log_dict.update({self._get_wandb_step_name(): step})
+        else:
+            log_dict.update({self._get_wandb_step_name(): self.global_step})
+
+        wandb.log(log_dict)
+
+    def save(self, milestone: str):
+        state = self.get_state()
+        torch.save(state, str(self.output_folder / f"model-{milestone}.pt"))
+
+    def load(self, milestone: str):
+        state = torch.load(
+            str(self.output_folder / f"model-{milestone}.pt"),
+            map_location=self.device,
+        )
+        self.load_state(state)
+
+    @abstractmethod
+    def train(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_model_state(self, train: bool = True):
+        raise NotImplementedError
+
+    @torch.inference_mode()
+    @abstractmethod
+    def eval(self, dataloader: torch.utils.data.DataLoader) -> Dict[str, object]:
+        raise NotImplementedError
+
+    @torch.inference_mode()
+    @abstractmethod
+    def eval_during_training(self) -> Dict[str, object]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def device(self) -> torch.DeviceObjType:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def global_step(self) -> int:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_state(self) -> Dict[str, object]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_state(self, state: Dict[str, object]):
+        raise NotImplementedError

nntool/utils/__init__.py

@@ -0,0 +1,6 @@
+# this is for backward compatibility
+from ..experiment.utils import (
+    get_current_time,
+    get_output_path,
+    read_toml_file,
+)

nntool/wandb/__init__.py

@@ -0,0 +1 @@
+from .config import WandbConfig, init_wandb

nntool/wandb/config.py

@@ -0,0 +1,118 @@
+import os
+import git
+import wandb
+import tomli
+import warnings
+
+from dataclasses import dataclass, field
+
+from wandb.sdk.lib.disabled import RunDisabled
+from wandb.sdk.wandb_run import Run
+
+
+@dataclass
+class WandbConfig:
+    """
+    Configuration class for Weights and Biases (wandb) integration.
+
+    :param api_key_config_file: The file path to the configuration file containing the wandb API key.
+                                The file should be a toml file with a `[wandb]` section. Default is an empty string.
+
+    :param project: The name of the project in wandb. Default is an empty string.
+
+    :param entity: The wandb user or team name. Default is an empty string.
+
+    :param name: The name of the wandb run. Leave blank to use the default run name. Default is an empty string.
+
+    :param notes: Notes or comments for the wandb run. Default is an empty string.
+
+    :param log_git_hash: Whether to log the current Git hash. Default is True.
+
+    :param log_code: Whether to log the current codebase. Default is True.
+
+    :param code_root: The root directory of the codebase to be logged. Default is the current directory (`.`).
+
+    :param code_ext: A list of file extensions for the code files to be logged. Default includes `.py` and `.sh`.
+    """
+
+    # wandb api key (toml file with [wandb] key field)
+    api_key_config_file: str = ""
+
+    # project name in wandb
+    project: str = ""
+
+    # wandb user name
+    entity: str = ""
+
+    # wandb run name (leave blacnk to use default name)
+    name: str = ""
+
+    # wandb run notes
+    notes: str = ""
+
+    # log git hash
+    log_git_hash: bool = True
+
+    # log code
+    log_code: bool = True
+
+    # code root
+    code_root: str = "."
+
+    # code file extensions
+    code_ext: list[str] = field(default_factory=lambda: [".py", ".sh"])
+
+
+def is_wandb_enabled():
+    return wandb.run is not None
+
+
+def init_wandb(args: WandbConfig, run_config: dict) -> Run | RunDisabled | None:
+    """initialize wandb and log the configuration
+
+    :param args: WandbConfig object
+    :param run_config: configuration dictionary to be logged
+    :return: wandb run object
+    """
+    project, entity = args.project, args.entity
+    if "WANDB_API_KEY" in os.environ:
+        warnings.warn("WANDB_API_KEY is found in environment variables. Using it.")
+        wandb.login(key=os.environ["WANDB_API_KEY"])
+    elif args.api_key_config_file:
+        with open(args.api_key_config_file, "rb") as config_file:
+            config_data = tomli.load(config_file)
+        wandb.login(key=config_data["wandb"]["key"])
+        project = config_data["wandb"].get("project", args.project)
+        entity = config_data["wandb"].get("entity", args.entity)
+    else:
+        warnings.warn(
+            "WANDB_API_KEY is not found in environment variables or the local key file."
+        )
+        pass
+
+    if args.log_git_hash:
+        repo = git.Repo(search_parent_directories=True)
+        git_hash = repo.head.object.hexsha
+        args.notes = (
+            f"{args.notes + ', ' if args.notes else ''}" + f"git hash: {git_hash}"
+        )
+
+    wandb_run = wandb.init(
+        project=project,
+        entity=entity,
+        name=args.name if args.name else None,
+        notes=args.notes,
+        config=run_config,
+    )
+    if not is_wandb_enabled():
+        warnings.warn(
+            "wandb is not enabled after intialization. Please check `wandb enabled`."
+        )
+
+    wandb_run.log_code(
+        root=args.code_root,
+        include_fn=lambda path, root: any(
+            [path.endswith(ext) for ext in args.code_ext]
+        ),
+    )
+    return wandb_run

nntool-1.3.0.dist-info/METADATA

@@ -0,0 +1,25 @@
+Metadata-Version: 2.1
+Name: nntool
+Version: 1.3.0
+Summary: neural network tool for research
+Home-page: https://github.com/jhliu17/nntool
+Author: Junhao Liu
+Author-email: junhaoliu17@gmail.com
+License: MIT
+Keywords: deep learning,neural network,research
+Requires-Python: >=3.9
+Requires-Dist: setuptools >=68.0.0
+Requires-Dist: submitit >=1.5.2
+Requires-Dist: tyro >=0.8.12
+Requires-Dist: matplotlib >=3.8.0
+Requires-Dist: seaborn >=0.13.2
+Requires-Dist: wandb >=0.15.0
+Requires-Dist: tomli >=2.0.1
+Requires-Dist: cythonpackage
+Provides-Extra: dev
+Requires-Dist: pytest >=8.0.2 ; extra == 'dev'
+Requires-Dist: jax[cpu] >=0.4.0 ; extra == 'dev'
+Requires-Dist: torch >=2.2.0 ; extra == 'dev'
+Requires-Dist: mypy ; extra == 'dev'
+Requires-Dist: accelerate ; extra == 'dev'
+

nntool-1.3.0.dist-info/RECORD

@@ -0,0 +1,76 @@
+nntool/__init__.py,sha256=aLOwJwXql9j87YPh8ACPWuGMD4Z_R8JlJVh2P5ZdohY,75
+nntool/slurm/config.py,sha256=UFleYQuXLxz5F4Hx5WuN832UnA0uTQ05M_YyTvsnFOs,6860
+nntool/slurm/function.py,sha256=aSbAgNvYPZaWJoxgjHomVaIZMSRrRzV5XIz1w6Hnrqs,8163
+nntool/slurm/wrap.py,sha256=dN2lzRfE32WRaqj4ohUp9K4n8EKkk657zOer6FbRGPU,9044
+nntool/slurm/__init__.py,sha256=MHg76NOj0qEzr_sXiL0nTarVm6xIXymtbmVc0U9E4_o,228
+nntool/slurm/task.py,sha256=djxm5Ggi2EvPhijDNPCvx-Ppk5_BKQpTy1BAKPZhVHI,7580
+nntool/slurm/accelerator/utils.py,sha256=oZrpcvlHYH14j-uNVEiSye6KiuULG1E34zGLw_IczHc,1298
+nntool/slurm/accelerator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nntool/slurm/accelerator/__pycache__/utils.cpython-310.pyc,sha256=o-bxDxmitSWbG6o0KuLgTRXZLe3AHVPRoGzq1Ev-iTU,1435
+nntool/slurm/accelerator/__pycache__/utils.cpython-310.opt-1.pyc,sha256=o-bxDxmitSWbG6o0KuLgTRXZLe3AHVPRoGzq1Ev-iTU,1435
+nntool/slurm/accelerator/__pycache__/__init__.cpython-310.pyc,sha256=qmfQ9skriIu7ee28WpFjf5n6BjOLu0XSvZT-fiRe3VU,129
+nntool/slurm/accelerator/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=qmfQ9skriIu7ee28WpFjf5n6BjOLu0XSvZT-fiRe3VU,129
+nntool/slurm/csrc/__init__.py,sha256=ew8vy7OKCavl2vOb0ZCDKunHjNGYsfgqZvse7UyhjkE,104
+nntool/slurm/csrc/__compile__.cpython-310-x86_64-linux-gnu.so,sha256=oDiT0iKdYo6SnwClCh55g6H8G1NE-77C9Xmfn6ZiURA,1813928
+nntool/slurm/csrc/__pycache__/__init__.cpython-310.pyc,sha256=eT1nHcuPo6fWx7d_BY-pSr46zN-sR4hr0Q6uS7FoBIY,240
+nntool/slurm/csrc/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=pE9vEMhMyI_hM8OSnP6nRgJUqHCFsfO5aiX6Fxofft8,240
+nntool/slurm/__pycache__/function.cpython-310.pyc,sha256=CQ7NiikaZGjLS-670_yOobPLZjluUwo4_iFopouGV4U,7972
+nntool/slurm/__pycache__/task.cpython-310.opt-1.pyc,sha256=a_zpjcS1-LNYbzvQmkUmknWNuC7FfsdrrYVyL116mxE,7634
+nntool/slurm/__pycache__/wrap.cpython-310.pyc,sha256=1O8zgL9Qrs7EyV6d1XjNWAu4XQqLQNO3NC499q2aCuw,8466
+nntool/slurm/__pycache__/config.cpython-310.opt-1.pyc,sha256=dJEeXTN48zdGN2-A1OMqDRgZCrd-e0FJ5NyQJ9UzBIs,5702
+nntool/slurm/__pycache__/task.cpython-310.pyc,sha256=dllCJK41Fw6ZsyvkVvG-FSemodTU9650FRxTVzrzNCk,7634
+nntool/slurm/__pycache__/wrap.cpython-310.opt-1.pyc,sha256=1O8zgL9Qrs7EyV6d1XjNWAu4XQqLQNO3NC499q2aCuw,8466
+nntool/slurm/__pycache__/function.cpython-310.opt-1.pyc,sha256=CQ7NiikaZGjLS-670_yOobPLZjluUwo4_iFopouGV4U,7972
+nntool/slurm/__pycache__/__init__.cpython-310.pyc,sha256=ftpFxoxqT4fcl9cyh7JFHRBBHm4YJ_RS95qJTBsfU7Y,400
+nntool/slurm/__pycache__/config.cpython-310.pyc,sha256=SeLnj_y4WyVJRO5UQ1KrzfLmC-xWt2UzIeW1d9GLWK4,5702
+nntool/slurm/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=ftpFxoxqT4fcl9cyh7JFHRBBHm4YJ_RS95qJTBsfU7Y,400
+nntool/train/trainer.py,sha256=f_cFFTdddT1EDX-7DN1_y-jQ3WWxO3ZoHXSF3rgBKqY,2808
+nntool/train/__init__.py,sha256=LRc7-qwMItUFM4FvB0fipqIRJfE_UV0SDBmGBUnMTGQ,827
+nntool/train/__pycache__/trainer.cpython-310.pyc,sha256=IZqX7G_cREw3vcoyoUx5LL9tW8Cdg-UCbmq0FO38hYw,3558
+nntool/train/__pycache__/trainer.cpython-310.opt-1.pyc,sha256=IZqX7G_cREw3vcoyoUx5LL9tW8Cdg-UCbmq0FO38hYw,3558
+nntool/train/__pycache__/__init__.cpython-310.pyc,sha256=K4kARzeBWVh2C_A8KLxNSZyE_bLzpd3NsS4wyCyaglU,582
+nntool/train/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=K4kARzeBWVh2C_A8KLxNSZyE_bLzpd3NsS4wyCyaglU,582
+nntool/utils/__init__.py,sha256=wo1H6JzA6F9VndjieS1UypPf78m2_9Io3eNqc336oxA,135
+nntool/utils/__pycache__/__init__.cpython-310.pyc,sha256=rjTPcYHdduVBookQAd7iDJUA7d7217sSx-yZD6WjtdI,225
+nntool/utils/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=rjTPcYHdduVBookQAd7iDJUA7d7217sSx-yZD6WjtdI,225
+nntool/experiment/utils.py,sha256=ajrjYNltyNwNfmFciygCvbcRBtn8Gc7-dVR07Q8xbKo,2119
+nntool/experiment/config.py,sha256=3PQx2B9bTVxdj0i2BQ_shHSzDdFn6nNfjUAfuVjVJoo,3861
+nntool/experiment/__init__.py,sha256=vd7SV9-29tceUk0QiWSgxYJIJws0M2lSSd2pu5KSKwE,127
+nntool/experiment/__pycache__/utils.cpython-310.pyc,sha256=B9zRohwMx_kyVq4i3ijfi20GBoPENmF2ipnajE2TY8I,1932
+nntool/experiment/__pycache__/config.cpython-310.opt-1.pyc,sha256=iS1ZQmJitzdE4Mf83j1pfmH3vtrGUnfFzeD5B4TuksU,4088
+nntool/experiment/__pycache__/utils.cpython-310.opt-1.pyc,sha256=B9zRohwMx_kyVq4i3ijfi20GBoPENmF2ipnajE2TY8I,1932
+nntool/experiment/__pycache__/__init__.cpython-310.pyc,sha256=VeRcGKZKmw3v7_BaWCLZ_idqmV00lYsmibu1gzEUAr4,270
+nntool/experiment/__pycache__/config.cpython-310.pyc,sha256=iS1ZQmJitzdE4Mf83j1pfmH3vtrGUnfFzeD5B4TuksU,4088
+nntool/experiment/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=VeRcGKZKmw3v7_BaWCLZ_idqmV00lYsmibu1gzEUAr4,270
+nntool/plot/__init__.py,sha256=M0qnyO48RzI_62v1EzhXNIszLZEJ7tsTS2YuzDoOoKE,133
+nntool/plot/context.py,sha256=9aEFWZSwliKCX5n2GX1uWPPJR_7dbn8XSVNefipkGB8,1331
+nntool/plot/csrc/__init__.py,sha256=fkzL5CkVtu6JFZ8xxMdLYjRt86bqWiHYfh3UI4iMi8s,51
+nntool/plot/csrc/__compile__.cpython-310-x86_64-linux-gnu.so,sha256=2GRlCicARxA4V-iuOdXuGu0deqT27CJXjk3y4xm0ErA,738224
+nntool/plot/csrc/__pycache__/__init__.cpython-310.pyc,sha256=9vJKTmP_6UeL9MkVQ4ZUlupDBosXvLlmuDaJHrrW9x8,174
+nntool/plot/csrc/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=iXVFNvhZVFai5AK9LShFIUDSYiPEKn0Umq15x1CpHMQ,174
+nntool/plot/__pycache__/context.cpython-310.opt-1.pyc,sha256=J_kOkmIV_GQPBuJ_hWQ7YYgGhuFjQpnvWbPRjJzEpD0,1750
+nntool/plot/__pycache__/__init__.cpython-310.pyc,sha256=GbhpEuu-hHu2B0A3A6dOwoev-mNveC01qEhJjXWvT3U,275
+nntool/plot/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=pS8ucjrevYJZOB2nBbAKXVi5uaizzCZV4NY-0pk053E,275
+nntool/plot/__pycache__/context.cpython-310.pyc,sha256=J_kOkmIV_GQPBuJ_hWQ7YYgGhuFjQpnvWbPRjJzEpD0,1750
+nntool/parser/__init__.py,sha256=991hqyHOJjoaN5WDrIBhEduZs9rmvgV6ShUeZNDAiXo,34
+nntool/parser/parse.py,sha256=IqlcqQfd7GoeTGJkMOrb2_-x0s49Zz_2fA3Fsk_6ysQ,552
+nntool/parser/__pycache__/parse.cpython-310.opt-1.pyc,sha256=v7j8kcFMlDt6-p7cHzgBXRP-vJxEiq4ZFCVgrfw4VFs,601
+nntool/parser/__pycache__/__init__.cpython-310.pyc,sha256=7pZwE5lGf6dnEJsprE9TqUlrKUkvrH9pC9h4WvLR2cw,162
+nntool/parser/__pycache__/parse.cpython-310.pyc,sha256=v7j8kcFMlDt6-p7cHzgBXRP-vJxEiq4ZFCVgrfw4VFs,601
+nntool/parser/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=7pZwE5lGf6dnEJsprE9TqUlrKUkvrH9pC9h4WvLR2cw,162
+nntool/__pycache__/__init__.cpython-310.pyc,sha256=vUDWdiFPntXWhlcy1OzRgdLU7U9oeNV-eT0LlRvxIMA,277
+nntool/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=vUDWdiFPntXWhlcy1OzRgdLU7U9oeNV-eT0LlRvxIMA,277
+nntool/wandb/config.py,sha256=F395QOoByO1-f5ocAmNfEGa7UfZjAWghYxoQgNKV-Ks,3605
+nntool/wandb/__init__.py,sha256=lvn7f3wMkVRz0oyYSjfasFQKjmJ5d6wkDpIDgf0DDKs,44
+nntool/wandb/__pycache__/config.cpython-310.opt-1.pyc,sha256=6LZQaVixEDhAJuPNCh-5kOxc5z_kWBL2tEHWLzFufFc,3783
+nntool/wandb/__pycache__/__init__.cpython-310.pyc,sha256=RELrQA63lgAHAzHWXzstH2k3vmlBuoKXWYYoOeFGq90,180
+nntool/wandb/__pycache__/config.cpython-310.pyc,sha256=SpyfZu8KSSkRuAcpywRrKDiV_C5eNd6T8J3BTuhanZs,3783
+nntool/wandb/__pycache__/__init__.cpython-310.opt-1.pyc,sha256=RELrQA63lgAHAzHWXzstH2k3vmlBuoKXWYYoOeFGq90,180
+nntool-1.3.0.dist-info/METADATA,sha256=Rf-21UEryukOszH46JqHrcdg4C9cbwcljBIIHZFJXKk,769
+nntool-1.3.0.dist-info/._requires.txt,sha256=yNlZBj9Qbibg_4etwM6NkI38o7TjR6UoE0tEoNdYPjY,267
+nntool-1.3.0.dist-info/RECORD,,
+nntool-1.3.0.dist-info/WHEEL,sha256=CzQQWV-lNyM92gr3iaBk8dvO35YDHRxgzkZ-dxumUIM,152
+nntool-1.3.0.dist-info/top_level.txt,sha256=hXO7zSAI-o9plYeUkCiKbkDpQCfTDIfjivyoSEqT28Y,7
+nntool-1.3.0.dist-info/._SOURCES.txt,sha256=yNlZBj9Qbibg_4etwM6NkI38o7TjR6UoE0tEoNdYPjY,267
+nntool-1.3.0.dist-info/._top_level.txt,sha256=MTm3Nh1CVxvWnLnzanzf6CglfDcr-AzoerhtjweyVT4,366
+nntool-1.3.0.dist-info/._dependency_links.txt,sha256=MTm3Nh1CVxvWnLnzanzf6CglfDcr-AzoerhtjweyVT4,366

nntool-1.3.0.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: bdist_wheel (0.43.0)
+Root-Is-Purelib: false
+Tag: cp310-cp310-manylinux_2_17_x86_64
+Tag: cp310-cp310-manylinux2014_x86_64
+

nntool-1.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+nntool