nshtrainer 0.1.0__tar.gz

nshtrainer-0.1.0/PKG-INFO

@@ -0,0 +1,18 @@
+Metadata-Version: 2.1
+Name: nshtrainer
+Version: 0.1.0
+Summary: 
+Author: Nima Shoghi
+Author-email: nimashoghi@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
+Requires-Dist: nshrunner (>=0.1.0,<0.2.0)
+Requires-Dist: torch
+Requires-Dist: typing-extensions
+Description-Content-Type: text/markdown
+
+

nshtrainer-0.1.0/pyproject.toml

@@ -0,0 +1,18 @@
+[tool.poetry]
+name = "nshtrainer"
+version = "0.1.0"
+description = ""
+authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+nshconfig = "^0.2.0"
+nshrunner = "^0.1.0"
+torch = "*"
+typing-extensions = "*"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

nshtrainer-0.1.0/src/nshtrainer/__init__.py

@@ -0,0 +1,64 @@
+from . import _experimental as _experimental
+from . import actsave as actsave
+from . import callbacks as callbacks
+from . import lr_scheduler as lr_scheduler
+from . import model as model
+from . import nn as nn
+from . import optimizer as optimizer
+from . import snapshot as snapshot
+from . import typecheck as typecheck
+from ._snoop import snoop as snoop
+from .actsave import ActLoad as ActLoad
+from .actsave import ActSave as ActSave
+from .config import MISSING as MISSING
+from .config import AllowMissing as AllowMissing
+from .config import Field as Field
+from .config import MissingField as MissingField
+from .config import PrivateAttr as PrivateAttr
+from .config import TypedConfig as TypedConfig
+from .data import dataset_transform as dataset_transform
+from .log import init_python_logging as init_python_logging
+from .log import lovely as lovely
+from .log import pretty as pretty
+from .lr_scheduler import LRSchedulerConfig as LRSchedulerConfig
+from .model import ActSaveConfig as ActSaveConfig
+from .model import Base as Base
+from .model import BaseConfig as BaseConfig
+from .model import BaseLoggerConfig as BaseLoggerConfig
+from .model import BaseProfilerConfig as BaseProfilerConfig
+from .model import CheckpointLoadingConfig as CheckpointLoadingConfig
+from .model import CheckpointSavingConfig as CheckpointSavingConfig
+from .model import ConfigList as ConfigList
+from .model import DirectoryConfig as DirectoryConfig
+from .model import (
+    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
+)
+from .model import EnvironmentConfig as EnvironmentConfig
+from .model import (
+    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+)
+from .model import (
+    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
+)
+from .model import GradientClippingConfig as GradientClippingConfig
+from .model import LightningDataModuleBase as LightningDataModuleBase
+from .model import LightningModuleBase as LightningModuleBase
+from .model import LoggingConfig as LoggingConfig
+from .model import MetricConfig as MetricConfig
+from .model import OptimizationConfig as OptimizationConfig
+from .model import PrimaryMetricConfig as PrimaryMetricConfig
+from .model import PythonLogging as PythonLogging
+from .model import ReproducibilityConfig as ReproducibilityConfig
+from .model import RunnerConfig as RunnerConfig
+from .model import SanityCheckingConfig as SanityCheckingConfig
+from .model import SeedConfig as SeedConfig
+from .model import TrainerConfig as TrainerConfig
+from .model import WandbWatchConfig as WandbWatchConfig
+from .nn import TypedModuleDict as TypedModuleDict
+from .nn import TypedModuleList as TypedModuleList
+from .optimizer import OptimizerConfig as OptimizerConfig
+from .runner import Runner as Runner
+from .runner import SnapshotConfig as SnapshotConfig
+from .trainer import Trainer as Trainer
+from .util.singleton import Registry as Registry
+from .util.singleton import Singleton as Singleton

nshtrainer-0.1.0/src/nshtrainer/_experimental/__init__.py

@@ -0,0 +1,2 @@
+from .flops import MEASURE_FLOPS_AVAILABLE as MEASURE_FLOPS_AVAILABLE
+from .flops import measure_flops as measure_flops

nshtrainer-0.1.0/src/nshtrainer/_experimental/flops/__init__.py

@@ -0,0 +1,48 @@
+from collections.abc import Callable
+
+import torch
+from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_1
+
+MEASURE_FLOPS_AVAILABLE = _TORCH_GREATER_EQUAL_2_1
+
+
+def measure_flops(
+    forward_fn: Callable[[], torch.Tensor],
+    loss_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+    display: bool = True,
+) -> int:
+    """Utility to compute the total number of FLOPs used by a module during training or during inference.
+
+    It's recommended to create a meta-device model for this:
+
+    Example::
+
+        with torch.device("meta"):
+            model = MyModel()
+            x = torch.randn(2, 32)
+
+        model_fwd = lambda: model(x)
+        fwd_flops = measure_flops(model, model_fwd)
+
+        model_loss = lambda y: y.sum()
+        fwd_and_bwd_flops = measure_flops(model, model_fwd, model_loss)
+
+    Args:
+        model: The model whose FLOPs should be measured.
+        forward_fn: A function that runs ``forward`` on the model and returns the result.
+        loss_fn: A function that computes the loss given the ``forward_fn`` output. If provided, the loss and `backward`
+            FLOPs will be included in the result.
+
+    """
+    if not MEASURE_FLOPS_AVAILABLE:
+        raise ImportError("`measure_flops` requires PyTorch >= 2.1.")
+
+    from .flop_counter import FlopCounterMode
+
+    flop_counter = FlopCounterMode(display=display)
+    with flop_counter:
+        if loss_fn is None:
+            forward_fn()
+        else:
+            loss_fn(forward_fn()).backward()
+    return flop_counter.get_total_flops()