nntool 1.3.0__cp39-cp39-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-39.pyc,sha256=7doqmi9AEji7UyCnQ6X7ddu0H62IWmcezuqpNP5DlAE,1443
+nntool/slurm/accelerator/__pycache__/__init__.cpython-39.pyc,sha256=kaXilaiufXK9jRyNsYKJtqssUXJfzghbAlRZrLnhTzg,127
+nntool/slurm/accelerator/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=kaXilaiufXK9jRyNsYKJtqssUXJfzghbAlRZrLnhTzg,127
+nntool/slurm/accelerator/__pycache__/utils.cpython-39.opt-1.pyc,sha256=7doqmi9AEji7UyCnQ6X7ddu0H62IWmcezuqpNP5DlAE,1443
+nntool/slurm/csrc/__init__.py,sha256=ew8vy7OKCavl2vOb0ZCDKunHjNGYsfgqZvse7UyhjkE,104
+nntool/slurm/csrc/__compile__.cpython-39-x86_64-linux-gnu.so,sha256=mahoZNJP_0WoGrTSNpx7gtI4UhhwqQa6B5ix7ekatrM,1821168
+nntool/slurm/csrc/__pycache__/__init__.cpython-39.pyc,sha256=3E9i-bc-bm5XW5rNk-HLyfQjdFjvwjn50eBIOik4AzU,238
+nntool/slurm/csrc/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=aVJ3yJQPtRtwKwkLSKip_1WNzIyQ6BB1psH2SCAcVXw,238
+nntool/slurm/__pycache__/function.cpython-39.opt-1.pyc,sha256=HFR2rm67ba4K_7l9ilR08Q-n8WYa4RKAqDQYtLXbxXQ,7977
+nntool/slurm/__pycache__/config.cpython-39.pyc,sha256=ZT4rON8FmKXZ4ky5SsRcYAcehJDbOF_o6pFITWnPKzA,5709
+nntool/slurm/__pycache__/task.cpython-39.pyc,sha256=VD6JMOJuP25sUeP0Y3T2EEdwkMSTXtqYbRPhKWj9cxU,7605
+nntool/slurm/__pycache__/__init__.cpython-39.pyc,sha256=Eimcm8-iLgkvIdqxLTJd77CQKTwTFkGlws-aBqYZA4A,398
+nntool/slurm/__pycache__/config.cpython-39.opt-1.pyc,sha256=2Gyto7BQpNxTzhcF5ig9FNl16cKtuf_nXwGeNMtN5Do,5709
+nntool/slurm/__pycache__/wrap.cpython-39.opt-1.pyc,sha256=RKHEslEmz8OyFcnD7g1nQHlrapRTOjmRTg8aIQn1WQo,8426
+nntool/slurm/__pycache__/function.cpython-39.pyc,sha256=HFR2rm67ba4K_7l9ilR08Q-n8WYa4RKAqDQYtLXbxXQ,7977
+nntool/slurm/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=Eimcm8-iLgkvIdqxLTJd77CQKTwTFkGlws-aBqYZA4A,398
+nntool/slurm/__pycache__/wrap.cpython-39.pyc,sha256=RKHEslEmz8OyFcnD7g1nQHlrapRTOjmRTg8aIQn1WQo,8426
+nntool/slurm/__pycache__/task.cpython-39.opt-1.pyc,sha256=0PHFWliIblp_RKaoyUSdPOFSIzlRDxKsrM7nEiJNDXA,7605
+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-39.opt-1.pyc,sha256=Koa-n3AAw7hxUFhnCwrb2OeKj4ouo5yGcKu0pQC00Qk,3640
+nntool/train/__pycache__/__init__.cpython-39.pyc,sha256=XJ4avfImgGXls8MjcDMd1BBd6r9BvmM8l5DY9q49s4E,574
+nntool/train/__pycache__/trainer.cpython-39.pyc,sha256=Koa-n3AAw7hxUFhnCwrb2OeKj4ouo5yGcKu0pQC00Qk,3640
+nntool/train/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=XJ4avfImgGXls8MjcDMd1BBd6r9BvmM8l5DY9q49s4E,574
+nntool/utils/__init__.py,sha256=wo1H6JzA6F9VndjieS1UypPf78m2_9Io3eNqc336oxA,135
+nntool/utils/__pycache__/__init__.cpython-39.pyc,sha256=Ecr_FN-nIlxdTIQUfRxjsSrn86UFse8GL4OALEy3ntQ,223
+nntool/utils/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=Ecr_FN-nIlxdTIQUfRxjsSrn86UFse8GL4OALEy3ntQ,223
+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-39.pyc,sha256=c9YpeYqT6ut4C7iayFZR59_03JAEGKk1N-o4v-7NONU,1902
+nntool/experiment/__pycache__/config.cpython-39.pyc,sha256=ThrOSF1m-tZN0bUh2GNI9fFvq6b4CpVRILNApNwtWNE,4094
+nntool/experiment/__pycache__/__init__.cpython-39.pyc,sha256=vkbYb5UQf3qyq9hwd7mqnKPNi0PDi0d-c0bhQkwggAw,268
+nntool/experiment/__pycache__/config.cpython-39.opt-1.pyc,sha256=ThrOSF1m-tZN0bUh2GNI9fFvq6b4CpVRILNApNwtWNE,4094
+nntool/experiment/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=vkbYb5UQf3qyq9hwd7mqnKPNi0PDi0d-c0bhQkwggAw,268
+nntool/experiment/__pycache__/utils.cpython-39.opt-1.pyc,sha256=c9YpeYqT6ut4C7iayFZR59_03JAEGKk1N-o4v-7NONU,1902
+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-39-x86_64-linux-gnu.so,sha256=jWadcwppOzlQqyDxBD573bJgaWqTw_4Xnz9bNQtzKK4,743512
+nntool/plot/csrc/__pycache__/__init__.cpython-39.pyc,sha256=ugIdDHIbyqmBhiQRJP0cDGtKuU8UP8Ee4_mBYXtGPyY,172
+nntool/plot/csrc/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=Q37daVDg4xzWBtKaz01-P8M4ARQV7cON__9ObIsTdec,172
+nntool/plot/__pycache__/context.cpython-39.pyc,sha256=RrrD4ri3FAxutyN7JXln9a8fAeeTXItruYqrA-5xtOA,1728
+nntool/plot/__pycache__/context.cpython-39.opt-1.pyc,sha256=RrrD4ri3FAxutyN7JXln9a8fAeeTXItruYqrA-5xtOA,1728
+nntool/plot/__pycache__/__init__.cpython-39.pyc,sha256=4Smux6gn6pzTu8Jy_aCgZZpZBRcLp_pn1BGISpEgQik,273
+nntool/plot/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=bW7IMhsAzuFh1CVdS85VwZYN6hh3xZpwelEBcGXnDzQ,273
+nntool/parser/__init__.py,sha256=991hqyHOJjoaN5WDrIBhEduZs9rmvgV6ShUeZNDAiXo,34
+nntool/parser/parse.py,sha256=IqlcqQfd7GoeTGJkMOrb2_-x0s49Zz_2fA3Fsk_6ysQ,552
+nntool/parser/__pycache__/parse.cpython-39.pyc,sha256=aC1pt4od9UQxXVcsQCJYS314rZG4kGhbhRX5K1PMvMs,593
+nntool/parser/__pycache__/__init__.cpython-39.pyc,sha256=1JkUPTfRoiA_1Vw3MiMRV48_GhHerNZGgVU3Q4RHtqw,160
+nntool/parser/__pycache__/parse.cpython-39.opt-1.pyc,sha256=aC1pt4od9UQxXVcsQCJYS314rZG4kGhbhRX5K1PMvMs,593
+nntool/parser/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=1JkUPTfRoiA_1Vw3MiMRV48_GhHerNZGgVU3Q4RHtqw,160
+nntool/__pycache__/__init__.cpython-39.pyc,sha256=XUWItoTVmdmESTFajbmjBp-g9U_a4dT-KZsgdksdj4A,275
+nntool/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=XUWItoTVmdmESTFajbmjBp-g9U_a4dT-KZsgdksdj4A,275
+nntool/wandb/config.py,sha256=F395QOoByO1-f5ocAmNfEGa7UfZjAWghYxoQgNKV-Ks,3605
+nntool/wandb/__init__.py,sha256=lvn7f3wMkVRz0oyYSjfasFQKjmJ5d6wkDpIDgf0DDKs,44
+nntool/wandb/__pycache__/config.cpython-39.pyc,sha256=zvWnc7xLtON-lvLGW9kbdLLo0jfvtxjk8FG2KTSlGuI,3769
+nntool/wandb/__pycache__/__init__.cpython-39.pyc,sha256=ohHQMM2D2jsMQFCpdJG-GI08Qei0fkkk1XPef7D1yGA,178
+nntool/wandb/__pycache__/config.cpython-39.opt-1.pyc,sha256=Y4jG7eHmkrm8SZt1FXNs1cNx3M9wYF3E6e72ZzoY2tc,3769
+nntool/wandb/__pycache__/__init__.cpython-39.opt-1.pyc,sha256=ohHQMM2D2jsMQFCpdJG-GI08Qei0fkkk1XPef7D1yGA,178
+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=rY0Y6THYM7EImsHfF-zs67o8pQciAsMw9_YuSvftjrQ,148
+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: cp39-cp39-manylinux_2_17_x86_64
+Tag: cp39-cp39-manylinux2014_x86_64
+

nntool-1.3.0.dist-info/top_level.txt

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