nntool 1.3.0__cp312-cp312-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__/__init__.cpython-312.opt-1.pyc,sha256=q3NjQ7KgKaJSvVjECPyoPrzZMSAs11ixUQCq9i22F5s,133
+nntool/slurm/accelerator/__pycache__/__init__.cpython-312.pyc,sha256=q3NjQ7KgKaJSvVjECPyoPrzZMSAs11ixUQCq9i22F5s,133
+nntool/slurm/accelerator/__pycache__/utils.cpython-312.pyc,sha256=bfZeF_8nADjSjCRWWmjV9Wa-3p_u8097Wn6RXFR_5yI,2075
+nntool/slurm/accelerator/__pycache__/utils.cpython-312.opt-1.pyc,sha256=bfZeF_8nADjSjCRWWmjV9Wa-3p_u8097Wn6RXFR_5yI,2075
+nntool/slurm/csrc/__init__.py,sha256=ew8vy7OKCavl2vOb0ZCDKunHjNGYsfgqZvse7UyhjkE,104
+nntool/slurm/csrc/__compile__.cpython-312-x86_64-linux-gnu.so,sha256=vTq0JVlf1sw9sEtAod0ueTdwgR5a_oNLzO5mlPxC4V0,2049056
+nntool/slurm/csrc/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=r8nEkXwO350xgykDg3HU-g_wyS5N8IgUp20BlGPfGu8,288
+nntool/slurm/csrc/__pycache__/__init__.cpython-312.pyc,sha256=r8nEkXwO350xgykDg3HU-g_wyS5N8IgUp20BlGPfGu8,288
+nntool/slurm/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=MD9ftHFI70Elf66zViZZiHKHop57wTtcmUgkts68uV4,419
+nntool/slurm/__pycache__/__init__.cpython-312.pyc,sha256=MD9ftHFI70Elf66zViZZiHKHop57wTtcmUgkts68uV4,419
+nntool/slurm/__pycache__/function.cpython-312.pyc,sha256=sHtq5mgJiCL1t6cAYO_SQ_Bz9FNstPoiMLxHUiAUfSs,9435
+nntool/slurm/__pycache__/function.cpython-312.opt-1.pyc,sha256=sHtq5mgJiCL1t6cAYO_SQ_Bz9FNstPoiMLxHUiAUfSs,9435
+nntool/slurm/__pycache__/task.cpython-312.opt-1.pyc,sha256=F6PXFoWC0zHGFba3bFRcfYDu057bsz8Ym-T99q6ayXQ,12982
+nntool/slurm/__pycache__/config.cpython-312.opt-1.pyc,sha256=aq86Gkal89sIxhwiwlo9cl70CqQex1zNCqUS3I0GXoo,6739
+nntool/slurm/__pycache__/config.cpython-312.pyc,sha256=aq86Gkal89sIxhwiwlo9cl70CqQex1zNCqUS3I0GXoo,6739
+nntool/slurm/__pycache__/wrap.cpython-312.opt-1.pyc,sha256=8KD80OwPVKmCEL3SknNfuLh4TIkinist3JL6hldhrwI,9737
+nntool/slurm/__pycache__/wrap.cpython-312.pyc,sha256=8KD80OwPVKmCEL3SknNfuLh4TIkinist3JL6hldhrwI,9737
+nntool/slurm/__pycache__/task.cpython-312.pyc,sha256=F6PXFoWC0zHGFba3bFRcfYDu057bsz8Ym-T99q6ayXQ,12982
+nntool/train/trainer.py,sha256=f_cFFTdddT1EDX-7DN1_y-jQ3WWxO3ZoHXSF3rgBKqY,2808
+nntool/train/__init__.py,sha256=LRc7-qwMItUFM4FvB0fipqIRJfE_UV0SDBmGBUnMTGQ,827
+nntool/train/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=RanaqCgdJr1c1aNW9CxMSmoyMhqS-KF6HORHAF9KFLA,795
+nntool/train/__pycache__/trainer.cpython-312.pyc,sha256=GRpOv7ygqd9mjvW_aaSJRkEpep4lXTezFbBEDRz62Ac,5399
+nntool/train/__pycache__/__init__.cpython-312.pyc,sha256=RanaqCgdJr1c1aNW9CxMSmoyMhqS-KF6HORHAF9KFLA,795
+nntool/train/__pycache__/trainer.cpython-312.opt-1.pyc,sha256=GRpOv7ygqd9mjvW_aaSJRkEpep4lXTezFbBEDRz62Ac,5399
+nntool/utils/__init__.py,sha256=wo1H6JzA6F9VndjieS1UypPf78m2_9Io3eNqc336oxA,135
+nntool/utils/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=otQrTKM72cOTZ4_2lZZeKcEd_hchaBJZgEOaLp6608s,239
+nntool/utils/__pycache__/__init__.cpython-312.pyc,sha256=otQrTKM72cOTZ4_2lZZeKcEd_hchaBJZgEOaLp6608s,239
+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__/__init__.cpython-312.opt-1.pyc,sha256=PAMfLRKpa40-i9KBDEkuiDsn8D_waIcA9noUg4IDL2g,285
+nntool/experiment/__pycache__/__init__.cpython-312.pyc,sha256=PAMfLRKpa40-i9KBDEkuiDsn8D_waIcA9noUg4IDL2g,285
+nntool/experiment/__pycache__/config.cpython-312.opt-1.pyc,sha256=AKiGDfWvUhiC7oHsSou-0DXGo31uJqeQGTgXxTKXRkQ,5283
+nntool/experiment/__pycache__/config.cpython-312.pyc,sha256=AKiGDfWvUhiC7oHsSou-0DXGo31uJqeQGTgXxTKXRkQ,5283
+nntool/experiment/__pycache__/utils.cpython-312.pyc,sha256=ae4d0NVCIgcG8t3gKNglzdcZVN4HVrSfaB7htHQAMIE,2688
+nntool/experiment/__pycache__/utils.cpython-312.opt-1.pyc,sha256=ae4d0NVCIgcG8t3gKNglzdcZVN4HVrSfaB7htHQAMIE,2688
+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-312-x86_64-linux-gnu.so,sha256=2SP2MXiuYsajGK65ccGF0cNoNx5qKh1I8kVUtYLvgMU,963832
+nntool/plot/csrc/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=orDr5-IshSJtRpeaHy4qey3j9tBdsDme8BwT3YyoMZA,221
+nntool/plot/csrc/__pycache__/__init__.cpython-312.pyc,sha256=orDr5-IshSJtRpeaHy4qey3j9tBdsDme8BwT3YyoMZA,221
+nntool/plot/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=ieZewauncybTnxkCxm1YwDZnzSR1BIOnoMO6eOIWw8k,329
+nntool/plot/__pycache__/__init__.cpython-312.pyc,sha256=ieZewauncybTnxkCxm1YwDZnzSR1BIOnoMO6eOIWw8k,329
+nntool/plot/__pycache__/context.cpython-312.pyc,sha256=VyyXd7y1s7Vb3NreoGM0-_Dggcffx_wSOyvbRrkr0a8,2414
+nntool/plot/__pycache__/context.cpython-312.opt-1.pyc,sha256=VyyXd7y1s7Vb3NreoGM0-_Dggcffx_wSOyvbRrkr0a8,2414
+nntool/parser/__init__.py,sha256=991hqyHOJjoaN5WDrIBhEduZs9rmvgV6ShUeZNDAiXo,34
+nntool/parser/parse.py,sha256=IqlcqQfd7GoeTGJkMOrb2_-x0s49Zz_2fA3Fsk_6ysQ,552
+nntool/parser/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=xBHUleFqh4cXtBb_6Lb0jRgu79DFszPDf-xMLkUMSD0,169
+nntool/parser/__pycache__/__init__.cpython-312.pyc,sha256=xBHUleFqh4cXtBb_6Lb0jRgu79DFszPDf-xMLkUMSD0,169
+nntool/parser/__pycache__/parse.cpython-312.pyc,sha256=B00bHyrPJxPxt1gMFvujG34LwIQI4ssGbBlls6W3rmw,782
+nntool/parser/__pycache__/parse.cpython-312.opt-1.pyc,sha256=B00bHyrPJxPxt1gMFvujG34LwIQI4ssGbBlls6W3rmw,782
+nntool/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=zFHghtpKEYgzZjqbu2XpP2XEo5qZPUkRTlahPqSQs5g,320
+nntool/__pycache__/__init__.cpython-312.pyc,sha256=zFHghtpKEYgzZjqbu2XpP2XEo5qZPUkRTlahPqSQs5g,320
+nntool/wandb/config.py,sha256=F395QOoByO1-f5ocAmNfEGa7UfZjAWghYxoQgNKV-Ks,3605
+nntool/wandb/__init__.py,sha256=lvn7f3wMkVRz0oyYSjfasFQKjmJ5d6wkDpIDgf0DDKs,44
+nntool/wandb/__pycache__/__init__.cpython-312.opt-1.pyc,sha256=7upnkW-1p6wAt02lGDuav27SjW4PaPZRVnoMsGlm5h0,190
+nntool/wandb/__pycache__/__init__.cpython-312.pyc,sha256=7upnkW-1p6wAt02lGDuav27SjW4PaPZRVnoMsGlm5h0,190
+nntool/wandb/__pycache__/config.cpython-312.opt-1.pyc,sha256=yduxm2GRHCwckEB1tq2S5cGF8ZouNCrmPs2n6bFCKi4,5124
+nntool/wandb/__pycache__/config.cpython-312.pyc,sha256=yduxm2GRHCwckEB1tq2S5cGF8ZouNCrmPs2n6bFCKi4,5124
+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=OvtvnbpcaxHa5TgXgfC48E6JA7zLr7svMspPP7Vk5o8,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: cp312-cp312-manylinux_2_17_x86_64
+Tag: cp312-cp312-manylinux2014_x86_64
+

nntool-1.3.0.dist-info/top_level.txt

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