nshtrainer 0.14.1__tar.gz → 0.14.4__tar.gz

{nshtrainer-0.14.1 → nshtrainer-0.14.4}/PKG-INFO

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

{nshtrainer-0.14.1 → nshtrainer-0.14.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.14.1"
+version = "0.14.4"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

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

@@ -0,0 +1 @@
+from lightning.fabric.utilities.throughput import measure_flops as measure_flops

{nshtrainer-0.14.1 → nshtrainer-0.14.4}/src/nshtrainer/runner.py

@@ -3,10 +3,9 @@ import functools
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from typing import Generic
 
-from nshrunner import RunInfo
+from nshrunner import RunInfo, Snapshot
 from nshrunner import Runner as _Runner
 from nshrunner._submit import screen
-from nshrunner.snapshot import SnapshotArgType
 from typing_extensions import TypeVar, TypeVarTuple, Unpack, override
 
 from .model.config import BaseConfig
@@ -93,7 +92,7 @@ class Runner(
         options: screen.ScreenJobKwargs = {},
         n_batches: int = 1,
         *,
-        snapshot: SnapshotArgType,
+        snapshot: Snapshot,
         setup_commands: Sequence[str] | None = None,
         env: Mapping[str, str] | None = None,
         transforms: list[

{nshtrainer-0.14.1 → nshtrainer-0.14.4}/src/nshtrainer/trainer/signal_connector.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import platform
 import re
 import signal
 import subprocess
@@ -25,6 +26,7 @@ log = logging.getLogger(__name__)
 
 _SIGNUM = int | signal.Signals
 _HANDLER: TypeAlias = Callable[[_SIGNUM, FrameType], Any] | int | signal.Handlers | None
+_IS_WINDOWS = platform.system() == "Windows"
 
 
 def _resolve_requeue_signals():
@@ -57,7 +59,7 @@ class _SignalConnector(_LightningSignalConnector):
         handlers: list[_HANDLER],
         replace_existing: bool = False,
     ):
-        if self._is_on_windows():
+        if _IS_WINDOWS:
             log.info(
                 f"Signal {signum.name} has no handlers or is not supported on Windows."
             )

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

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

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

@@ -1,48 +0,0 @@
-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()