nshtrainer 0.10.7__py3-none-any.whl → 0.10.9__py3-none-any.whl

nshtrainer/callbacks/gradient_skipping.py

@@ -1,8 +1,8 @@
-from logging import getLogger
-from typing import Literal, Protocol, runtime_checkable
+import importlib.util
+import logging
+from typing import Any, Literal, Protocol, runtime_checkable
 
 import torch
-import torchmetrics
 from lightning.pytorch import Callback, LightningModule, Trainer
 from torch.optim import Optimizer
 from typing_extensions import override
@@ -10,23 +10,29 @@ from typing_extensions import override
 from .base import CallbackConfigBase
 from .norm_logging import compute_norm
 
-log = getLogger(__name__)
+log = logging.getLogger(__name__)
 
 
 @runtime_checkable
 class HasGradSkippedSteps(Protocol):
-    grad_skipped_steps: torchmetrics.SumMetric
+    grad_skipped_steps: Any
 
 
 class GradientSkipping(Callback):
     def __init__(self, config: "GradientSkippingConfig"):
-        super().__init__()
+        if importlib.util.find_spec("torchmetrics") is not None:
+            raise ImportError(
+                "To use the GradientSkipping callback, please install torchmetrics: pip install torchmetrics"
+            )
 
+        super().__init__()
         self.config = config
 
     @override
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
         if not isinstance(pl_module, HasGradSkippedSteps):
+            import torchmetrics  # type: ignore
+
             pl_module.grad_skipped_steps = torchmetrics.SumMetric()
 
     @override
@@ -47,12 +53,7 @@ class GradientSkipping(Callback):
         ):
             return
 
-        norm = compute_norm(
-            pl_module,
-            optimizer,
-            self.config.norm_type,
-            grad=True,
-        )
+        norm = compute_norm(pl_module, optimizer, self.config.norm_type, grad=True)
 
         # If the norm is NaN/Inf, we don't want to skip the step
         # beacuse AMP checks for NaN/Inf grads to adjust the loss scale.

nshtrainer/data/transform.py

@@ -22,7 +22,13 @@ def transform(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx):
@@ -52,7 +58,13 @@ def transform_with_index(
         deepcopy: Whether to deep copy each item before applying the transform.
     """
 
-    import wrapt
+    try:
+        import wrapt
+    except ImportError:
+        raise ImportError(
+            "wrapt is not installed. wrapt is required for the transform function."
+            "Please install it using 'pip install wrapt'"
+        )
 
     class _TransformedWithIndexDataset(wrapt.ObjectProxy):
         def __getitem__(self, idx: int):

nshtrainer/model/_environment.py

@@ -9,7 +9,6 @@ from datetime import timedelta
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
-import git
 import nshconfig as C
 import psutil
 import torch
@@ -618,6 +617,11 @@ class GitRepositoryConfig(C.Config):
 
     @classmethod
     def from_current_directory(cls):
+        try:
+            import git
+        except ImportError:
+            return cls()
+
         draft = cls.draft()
         try:
             repo = git.Repo(os.getcwd(), search_parent_directories=True)

nshtrainer/model/config.py

@@ -286,6 +286,10 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     offline: bool = False
     """Whether to run WandB in offline mode."""
 
+    def offline_(self):
+        self.offline = True
+        return self
+
     @override
     def create_logger(self, root_config):
         if not self.enabled:

{nshtrainer-0.10.7.dist-info → nshtrainer-0.10.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.10.7
+Version: 0.10.9
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.10.7.dist-info → nshtrainer-0.10.9.dist-info}/RECORD

@@ -12,7 +12,7 @@ nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,
 nshtrainer/callbacks/early_stopping.py,sha256=jriSU761wf_qTJ9Bos0D3h5aDvZHYpRqK62Ne8aWp5I,3768
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
 nshtrainer/callbacks/finite_checks.py,sha256=AO5fa51uANAjAkeJfTquOjK6W_4RSU5Kky3f5jmAPlQ,2084
-nshtrainer/callbacks/gradient_skipping.py,sha256=fSJpjgHbztFKz7w3qFuCHZpmbEt9BCLAy-sU0B4xJQI,3474
+nshtrainer/callbacks/gradient_skipping.py,sha256=pqu5AELx4ctJxR2Y7YSSiGd5oGauVCTZFCEIIS6s88w,3665
 nshtrainer/callbacks/interval.py,sha256=smz5Zl8cN6X6yHKVsMRS2e3SEkzRCP3LvwE1ONvLfaw,8080
 nshtrainer/callbacks/latest_epoch_checkpoint.py,sha256=zCRAUsqW-2PaoIwVKlXOqdh2uF_B_YUUTmQO1wSomR8,2489
 nshtrainer/callbacks/log_epoch.py,sha256=fTa_K_Y8A7g09630cG4YkDE6AzSMPkjb9bpPm4gtqos,1120
@@ -25,7 +25,7 @@ nshtrainer/callbacks/timer.py,sha256=quS79oYClDUvQxJkNWmDMe0hwRUkkREgTgqzVrnom50
 nshtrainer/callbacks/wandb_watch.py,sha256=0JKDEPSDmiKverFm00lFfufOvtvr49akFXScvdUQnqc,2930
 nshtrainer/data/__init__.py,sha256=7mk1tr7SWUZ7ySbsf0y0ZPszk7u4QznPhQ-7wnpH9ec,149
 nshtrainer/data/balanced_batch_sampler.py,sha256=bcJBcQjh1hB1yKF_xSlT9AtEWv0BJjYc1CuH2BF-ea8,4392
-nshtrainer/data/transform.py,sha256=JeGxvytQly8hougrsdMmKG8gJ6qvFPDglJCO4Tp6STk,1795
+nshtrainer/data/transform.py,sha256=6SNs3_TpNpfhcwTwvPKyEJ3opM1OT7LmMEYQNHKgRl8,2227
 nshtrainer/ll/__init__.py,sha256=dD0ISxHJ2lg1HLSM0b3db7TBlsPpQCtChnuYO-c2oqI,2635
 nshtrainer/ll/_experimental.py,sha256=oBQCKOEVYoxuUU9eLb-Fg2B2mzZD7SA0zfAO6lmWZ88,53
 nshtrainer/ll/actsave.py,sha256=2lbiseSrjcwFT6AiyLNWarTWl1bnzliVWlu1iOfnP30,209
@@ -50,9 +50,9 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=hWWS4IXENRyH3RmJ7z1Wx1n3Lt1sNMlGOrcU6PW15o0,1104
 nshtrainer/model/__init__.py,sha256=TbexTxiE20WHYg5q3L88Hysk4LlHeKk_isv33aSBREA,1918
-nshtrainer/model/_environment.py,sha256=JCFxxwMhkviiMDkqIXJmiuepqiSYIlcoSQM7Y2H2KX4,23036
+nshtrainer/model/_environment.py,sha256=oTtecQeF5oY2RV7UkkSLnzDy3clz4AUkf9oocD6-e54,23115
 nshtrainer/model/base.py,sha256=Bmw-t70TydDbE9P0ee-lTibGoUhrCx5Qke-upa7FGVM,17512
-nshtrainer/model/config.py,sha256=6pAqDUk1eBloR3vZmtsWVdMrKeT2V3UvOn5UZ7YhZ_Q,53283
+nshtrainer/model/config.py,sha256=FRjn2pVbCzitBeIFcKXXPQ0TCmqvAwSylAvfedf1NHg,53356
 nshtrainer/model/modules/callback.py,sha256=JF59U9-CjJsAIspEhTJbVaGN0wGctZG7UquE3IS7R8A,6408
 nshtrainer/model/modules/debug.py,sha256=DTVty8cKnzj1GCULRyGx_sWTTsq9NLi30dzqjRTnuCU,1127
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -67,7 +67,6 @@ nshtrainer/nn/module_list.py,sha256=fb2u5Rqdjff8Pekyr9hkCPkBorQ-fldzzFAjsgWAm30,
 nshtrainer/nn/nonlinearity.py,sha256=owtU4kh4G98psD0axOJWVfBhm-OtJVgFM-TXSHmbNPU,3625
 nshtrainer/optimizer.py,sha256=kuJEA1pvB3y1FcsfhAoOJujVqEZqFHlmYO8GW6JeA1g,1527
 nshtrainer/runner.py,sha256=6qfE5FBONzD79kVHuWYKEvK0J_Qi5dMBbHQhRMmnIhE,3649
-nshtrainer/scripts/check_env.py,sha256=IMl6dSqsLYppI0XuCsVq8lK4bYqXwY9KHJkzsShz4Kg,806
 nshtrainer/scripts/find_packages.py,sha256=FbdlfmAefttFSMfaT0A46a-oHLP_ioaQKihwBfBeWeA,1467
 nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3E,40
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
@@ -79,6 +78,6 @@ nshtrainer/util/seed.py,sha256=HEXgVs-wldByahOysKwq7506OHxdYTEgmP-tDQVAEkQ,287
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.10.7.dist-info/METADATA,sha256=IQ6IEecsAvygnoV5P6_mkG9RjRGnb_cFuOf2Ic2HLIY,695
-nshtrainer-0.10.7.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.10.7.dist-info/RECORD,,
+nshtrainer-0.10.9.dist-info/METADATA,sha256=zpMnR1Jwc9kX7BhnHT6NA3nsPUwQzMAraBPOGsUSD1w,695
+nshtrainer-0.10.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.10.9.dist-info/RECORD,,

nshtrainer/scripts/check_env.py

@@ -1,41 +0,0 @@
-REQUIRED_PACKAGES = [
-    "beartype",
-    "cloudpickle",
-    "jaxtyping",
-    "lightning",
-    "lightning_fabric",
-    "lightning_utilities",
-    "lovely_numpy",
-    "lovely_tensors",
-    "numpy",
-    "psutil",
-    "pydantic",
-    "pydantic_core",
-    "pysnooper",
-    "rich",
-    "tabulate",
-    "torch",
-    "torchmetrics",
-    "tqdm",
-    "typing_extensions",
-    "wrapt",
-    "yaml",
-]
-
-
-def main():
-    import importlib.util
-    import sys
-
-    missing_packages: list[str] = []
-    for package_name in REQUIRED_PACKAGES:
-        spec = importlib.util.find_spec(package_name)
-        if spec is None:
-            missing_packages.append(package_name)
-
-    if missing_packages:
-        sys.exit(f"Error: Missing required packages: {', '.join(missing_packages)}")
-
-
-if __name__ == "__main__":
-    main()