nshtrainer 1.1.0__py3-none-any.whl → 1.1.1b1__py3-none-any.whl

nshtrainer/__init__.py

@@ -14,9 +14,6 @@ from .metrics import MetricConfig as MetricConfig
 from .model import LightningModuleBase as LightningModuleBase
 from .trainer import Trainer as Trainer
 from .trainer import TrainerConfig as TrainerConfig
-from .trainer import accelerator_registry as accelerator_registry
-from .trainer import callback_registry as callback_registry
-from .trainer import plugin_registry as plugin_registry
 
 try:
     from . import configs as configs

nshtrainer/configs/__init__.py

@@ -4,9 +4,6 @@ __codegen__ = True
 
 from nshtrainer import MetricConfig as MetricConfig
 from nshtrainer import TrainerConfig as TrainerConfig
-from nshtrainer import accelerator_registry as accelerator_registry
-from nshtrainer import callback_registry as callback_registry
-from nshtrainer import plugin_registry as plugin_registry
 from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
 from nshtrainer._directory import DirectoryConfig as DirectoryConfig
 from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
@@ -14,6 +11,7 @@ from nshtrainer._hf_hub import (
     HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
 )
 from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer._hf_hub import callback_registry as callback_registry
 from nshtrainer.callbacks import ActSaveConfig as ActSaveConfig
 from nshtrainer.callbacks import (
     BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
@@ -106,6 +104,8 @@ from nshtrainer.profiler import BaseProfilerConfig as BaseProfilerConfig
 from nshtrainer.profiler import ProfilerConfig as ProfilerConfig
 from nshtrainer.profiler import PyTorchProfilerConfig as PyTorchProfilerConfig
 from nshtrainer.profiler import SimpleProfilerConfig as SimpleProfilerConfig
+from nshtrainer.trainer import accelerator_registry as accelerator_registry
+from nshtrainer.trainer import plugin_registry as plugin_registry
 from nshtrainer.trainer._config import AcceleratorConfig as AcceleratorConfig
 from nshtrainer.trainer._config import (
     CheckpointCallbackConfig as CheckpointCallbackConfig,

nshtrainer/data/datamodule.py

@@ -30,9 +30,9 @@ class LightningDataModuleBase(
 
     @property
     @override
-    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+    def hparams_initial(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams = cast(THparams, super().hparams_initial)
-        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
+        return hparams
 
     @property
     @deprecated("Use `hparams` instead")

nshtrainer/model/base.py

@@ -134,16 +134,9 @@ class LightningModuleBase(
 
     @property
     @override
-    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+    def hparams_initial(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
         hparams = cast(THparams, super().hparams_initial)
-        hparams_dict = {"model": hparams.model_dump(mode="json")}
-        if (trainer := self._trainer) is not None:
-            from ..trainer import Trainer
-
-            if isinstance(trainer, Trainer):
-                hparams_dict["trainer"] = trainer.hparams.model_dump(mode="json")
-
-        return cast(Never, hparams_dict)
+        return hparams
 
     @property
     @deprecated("Use `hparams` instead")

nshtrainer/trainer/_log_hparams.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any, cast
+
+import nshconfig as C
+from lightning.pytorch import LightningDataModule, Trainer
+
+log = logging.getLogger(__name__)
+
+
+def _dict(obj: Any):
+    if isinstance(obj, C.Config):
+        return obj.model_dump(mode="json")
+
+    try:
+        return dict(obj)
+    except Exception:
+        return json.loads(
+            json.dumps(obj, default=lambda o: str(o), indent=4, sort_keys=True)
+        )
+
+
+def _dict_and_clean(obj: Any):
+    d = _dict(obj)
+
+    # Remove LightningCLI's internal hparam
+    d = {k: v for k, v in d.items() if k != "_class_path"}
+    return d
+
+
+def _log_hyperparams(trainer: Trainer) -> None:
+    if not trainer.loggers:
+        return
+
+    hparams_to_log: dict[str, Any] = {}
+
+    from .trainer import Trainer
+
+    if isinstance(trainer, Trainer):
+        hparams_to_log["trainer"] = _dict_and_clean(trainer.hparams)
+
+    if (
+        pl_module := trainer.lightning_module
+    ) is not None and pl_module._log_hyperparams:
+        hparams_to_log["model"] = _dict_and_clean(pl_module.hparams_initial)
+
+    if (
+        datamodule := cast(
+            LightningDataModule | None, getattr(trainer, "datamodule", None)
+        )
+    ) is not None and (datamodule._log_hyperparams):
+        hparams_to_log["datamodule"] = _dict_and_clean(datamodule.hparams_initial)
+
+    for logger in trainer.loggers:
+        logger.log_hyperparams(hparams_to_log)
+        logger.log_graph(pl_module)
+        logger.save()
+
+
+def patch_log_hparams_function():
+    try:
+        import lightning.pytorch.loggers.utilities
+        import lightning.pytorch.trainer.trainer
+
+        lightning.pytorch.loggers.utilities._log_hyperparams = _log_hyperparams
+        lightning.pytorch.trainer.trainer._log_hyperparams = _log_hyperparams
+        log.info(
+            "Patched lightning.pytorch's _log_hyperparams to use nshtrainer's version"
+        )
+    except ImportError:
+        pass
+
+    try:
+        import pytorch_lightning.loggers.utilities
+        import pytorch_lightning.trainer.trainer
+
+        pytorch_lightning.loggers.utilities._log_hyperparams = _log_hyperparams
+        pytorch_lightning.trainer.trainer._log_hyperparams = _log_hyperparams
+        log.info(
+            "Patched pytorch_lightning's _log_hyperparams to use nshtrainer's version"
+        )
+    except ImportError:
+        pass

nshtrainer/trainer/trainer.py

@@ -23,6 +23,7 @@ from ..callbacks.base import resolve_all_callbacks
 from ..util._environment_info import EnvironmentConfig
 from ..util.bf16 import is_bf16_supported_no_emulation
 from ._config import LightningTrainerKwargs, TrainerConfig
+from ._log_hparams import patch_log_hparams_function
 from ._runtime_callback import RuntimeTrackerCallback, Stage
 from .accelerator import AcceleratorConfigBase
 from .signal_connector import _SignalConnector
@@ -31,6 +32,9 @@ from .strategy import StrategyConfigBase
 log = logging.getLogger(__name__)
 
 
+patch_log_hparams_function()
+
+
 class Trainer(LightningTrainer):
     CHECKPOINT_HYPER_PARAMS_KEY = "trainer_hyper_parameters"
 

{nshtrainer-1.1.0.dist-info → nshtrainer-1.1.1b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.1.0
+Version: 1.1.1b1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.1.0.dist-info → nshtrainer-1.1.1b1.dist-info}/RECORD

@@ -1,5 +1,5 @@
 nshtrainer/.nshconfig.generated.json,sha256=yZd6cn1RhvNNJUgiUTRYut8ofZYvbulnpPG-rZIRhi4,106
-nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
+nshtrainer/__init__.py,sha256=VcqBfL8RgCcZDaY645nxeDmOspqerx4x46wggCMnS0E,692
 nshtrainer/_callback.py,sha256=ZDppiJ4d65tRXTEWYPZLH_F1xFizdz1pkWJe_sQ5uII,12564
 nshtrainer/_checkpoint/metadata.py,sha256=Hh5a7OkdknUEbkEwX6vS88-XLEeuVDoR6a3en2uLzQE,5597
 nshtrainer/_checkpoint/saver.py,sha256=utcrYKSosd04N9m2GIylufO5DO05D90qVU3mvadfApU,1658
@@ -32,7 +32,7 @@ nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU
 nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
-nshtrainer/configs/__init__.py,sha256=-rGk9pnRnuz4yKvACGOpY3nkrWnHholqZGk7UP2Vkrc,14716
+nshtrainer/configs/__init__.py,sha256=4WNs4Zv4PtHWD0KKH4X7j_zFt-COrEB0KhNIljsA6Rc,14740
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
@@ -101,7 +101,7 @@ nshtrainer/configs/util/config/dtype/__init__.py,sha256=PmGF-O4r6SXqEaagVsQ5YxEq
 nshtrainer/configs/util/config/duration/__init__.py,sha256=44lS2irOIPVfgshMTfnZM2jC6l0Pjst9w2M_lJoS_MU,353
 nshtrainer/data/__init__.py,sha256=K4i3Tw4g9EOK2zlMMbidi99y0SyI4w8P7_XUf1n42Ts,260
 nshtrainer/data/balanced_batch_sampler.py,sha256=r1cBKRXKHD8E1Ax6tj-FUbE-z1qpbO58mQ9VrK9uLnc,5481
-nshtrainer/data/datamodule.py,sha256=0M-HjGZQkLG77HXn4ZgLSypnbSjkjTq6GEJwGWe_gbM,4136
+nshtrainer/data/datamodule.py,sha256=Rb4-mA8iXtjRlNUHcIqVPEvxA_VkiJXwN1EvHIsydJ0,4095
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
 nshtrainer/loggers/__init__.py,sha256=fI0OHEltHP4tZI-KFB3npdzoxm_M2QsEYKxY3um05_s,592
 nshtrainer/loggers/actsave.py,sha256=wgNrpBB6wQM7qff8iLDb_sQnbiAcYHRmH56pcEJPB3o,1409
@@ -116,7 +116,7 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=irPyDjfUX843ze4bJM9sW8WSe
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=ox_ScK6V0J9nzIMhEB0qpToNKpt83VVgOVSRFCV-wBc,595
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=bZMNap0rkxRbAbu2BOHV_6YS2iZZnvy6wVSMOXGa_ZM,8680
+nshtrainer/model/base.py,sha256=LsOK5mMhYG5J0eSFKZKdd1fTvr38sgi8LLVSqoW6OCU,8386
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
@@ -134,6 +134,7 @@ nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
 nshtrainer/trainer/_config.py,sha256=s-_XoLc9mbNAdroRJyOKd3dLTyrFLQkPyGJkKDmBYf8,33267
+nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
 nshtrainer/trainer/plugin/__init__.py,sha256=q_q98MYNaZ2VE_tqGqYlQjQnlaF4NE1FUqVVbj0EK7k,517
@@ -144,7 +145,7 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=-BbEyWZ063O7tZme7Gdu1lVxK6p1NeuLc
 nshtrainer/trainer/plugin/precision.py,sha256=7lf7KZd_yFyPmhLApjEIv0pkoDB5zdxi-7in0wRj3z8,5436
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=Lo3vUo3ooTAjaX2fUYPFSMv5FP7sWfVov0QbA-T5hZ8,21113
+nshtrainer/trainer/trainer.py,sha256=BKRicDlLI7KstzuP0SmzJzp0U4GK5lhZcKHS1IuL5sA,21197
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
@@ -156,6 +157,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.1.0.dist-info/METADATA,sha256=9prl1vhkmagIgO0WmzfDEknzv4S0a0qDAh-Bj3LxFjI,960
-nshtrainer-1.1.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.1.0.dist-info/RECORD,,
+nshtrainer-1.1.1b1.dist-info/METADATA,sha256=wdOIQ91eUgWrIHfPLP06FD4uMkyyIfToR3VhBY-BXsE,962
+nshtrainer-1.1.1b1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.1.1b1.dist-info/RECORD,,