fix(measure): every measurement leaves the model in the mode it came in

fasterbench/_modidx.py

@@ -50,6 +50,7 @@
                                   'fasterbench.core._fmt_macs': ('core/core.html#_fmt_macs', 'fasterbench/core.py'),
                                   'fasterbench.core._fmt_table': ('core/core.html#_fmt_table', 'fasterbench/core.py'),
                                   'fasterbench.core._is_quantized': ('core/core.html#_is_quantized', 'fasterbench/core.py'),
+                                  'fasterbench.core._measuring': ('core/core.html#_measuring', 'fasterbench/core.py'),
                                   'fasterbench.core._on_device': ('core/core.html#_on_device', 'fasterbench/core.py'),
                                   'fasterbench.core._run_on_devices': ('core/core.html#_run_on_devices', 'fasterbench/core.py'),
                                   'fasterbench.core._section': ('core/core.html#_section', 'fasterbench/core.py'),

fasterbench/benchmark.py

@@ -126,6 +126,7 @@ def benchmark(
     if sample.dim() == 0:
         raise ValueError("sample must have at least 1 dimension")
 
+    # manual restore: also covers compute_compute, which never sets eval mode
     was_training = model.training
     model.eval()
 
@@ -138,5 +139,4 @@ def benchmark(
             energy=compute_energy_multi(model, sample, devices=energy_devices, **kwargs) if "energy" in metrics_set else {},
         )
     finally:
-        if was_training:
-            model.train()
+        model.train(was_training)

fasterbench/core.py

@@ -153,3 +153,14 @@ def _run_on_devices(
             warnings.warn(f"Unexpected error during {metric_name} benchmark on {d}: {e}")
             out[d_str] = nan_factory(d_str)
     return out
+
+# %% ../nbs/core/core.ipynb #44ce3eab
+@contextlib.contextmanager
+def _measuring(model: torch.nn.Module, dev: str | torch.device):
+    """Yield `model` on `dev` in eval mode, then restore its device and training mode."""
+    was_training = model.training
+    try:
+        with _on_device(model, dev) as m:
+            yield m.eval()
+    finally:
+        model.train(was_training)

fasterbench/energy.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from .core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported
+from .core import _device_ctx, _measuring, _sync, _run_on_devices, _ensure_device_supported
 
 try:
     from codecarbon import EmissionsTracker, OfflineEmissionsTracker
@@ -95,7 +95,7 @@ def compute_energy(
 
     with _device_ctx(device) as dev:
         _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-        with _on_device(model.eval(), dev) as model:
+        with _measuring(model, dev) as model:
             sample = sample.to(dev, non_blocking=True)
 
             for _ in range(warmup):

fasterbench/memory.py

@@ -6,7 +6,7 @@
 import warnings
 from dataclasses import dataclass, asdict
 from typing import Sequence
-from .core import _bytes_to_mib, _on_device, _run_on_devices, _ensure_device_supported
+from .core import _bytes_to_mib, _measuring, _run_on_devices, _ensure_device_supported
 
 import numpy as np
 import torch
@@ -50,7 +50,7 @@ def _gpu_metrics(
     """Measure GPU memory usage."""
     dev = torch.device("cuda")
     _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-    with _on_device(model.eval(), dev) as model:
+    with _measuring(model, dev) as model:
         sample = sample.to(dev)
 
         for _ in range(warmup):
@@ -88,7 +88,7 @@ def _cpu_metrics(
         return _nan_memory_metrics("cpu")
 
     proc = psutil.Process()
-    with _on_device(model.eval(), "cpu") as model:
+    with _measuring(model, "cpu") as model:
         sample = sample.cpu()
 
         rss0 = proc.memory_info().rss

fasterbench/profiling.py

@@ -11,7 +11,7 @@
 import torch
 import torch.nn as nn
 
-from .core import _bytes_to_mib, _device_ctx, _on_device, _section, _fmt_table, _fmt_float, _fmt_macs
+from .core import _bytes_to_mib, _device_ctx, _measuring, _section, _fmt_table, _fmt_float, _fmt_macs
 
 # %% auto #0
 __all__ = ['LayerProfiler']
@@ -177,7 +177,7 @@ def _profile_layers(
     config = _METRIC_CONFIG[metric]
     col = config["col"]
 
-    with _device_ctx(device) as dev, _on_device(model.eval(), dev) as model:
+    with _device_ctx(device) as dev, _measuring(model, dev) as model:
         sample = sample.to(dev)
 
         leaf_mods = _leaf_modules(model)

fasterbench/roofline.py

@@ -14,7 +14,7 @@
 import torch.nn as nn
 import plotly.graph_objects as go
 
-from .core import _device_ctx, _on_device, _sync, _fmt_float, _section
+from .core import _device_ctx, _measuring, _sync, _fmt_float, _section
 from .profiling import _leaf_modules, _tensor_bytes, _output_bytes
 
 # %% auto #0
@@ -304,7 +304,7 @@ def profile(
             if self.peaks is None:
                 self.peaks = measure_peaks(dev, dtype=self.sample.dtype)
 
-            with _on_device(self.model.eval(), dev) as model:
+            with _measuring(self.model, dev) as model:
                 sample = self.sample.to(dev)
                 leaf_mods = _leaf_modules(model)
 

fasterbench/speed.py

@@ -11,7 +11,7 @@
 import torch.nn as nn
 from torch.utils.benchmark import Timer
 
-from .core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported
+from .core import _device_ctx, _measuring, _sync, _run_on_devices, _ensure_device_supported
 
 # %% auto #0
 __all__ = ['SpeedMetrics', 'compute_speed', 'compute_speed_multi', 'sweep_threads', 'sweep_batch_sizes', 'sweep_latency']
@@ -69,7 +69,7 @@ def _forward_latencies(
 
     with _device_ctx(device) as dev:
         _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-        with _on_device(model.eval(), dev) as model:
+        with _measuring(model, dev) as model:
             sample = sample.to(dev, non_blocking=True)
 
             for _ in range(warmup):

nbs/analysis/benchmark.ipynb

@@ -175,6 +175,7 @@
     "    if sample.dim() == 0:\n",
     "        raise ValueError(\"sample must have at least 1 dimension\")\n",
     "\n",
+    "    # manual restore: also covers compute_compute, which never sets eval mode\n",
     "    was_training = model.training\n",
     "    model.eval()\n",
     "\n",
@@ -187,8 +188,7 @@
     "            energy=compute_energy_multi(model, sample, devices=energy_devices, **kwargs) if \"energy\" in metrics_set else {},\n",
     "        )\n",
     "    finally:\n",
-    "        if was_training:\n",
-    "            model.train()"
+    "        model.train(was_training)"
    ]
   },
   {

nbs/analysis/profiling.ipynb

@@ -46,7 +46,7 @@
     "import torch\n",
     "import torch.nn as nn\n",
     "\n",
-    "from fasterbench.core import _bytes_to_mib, _device_ctx, _on_device, _section, _fmt_table, _fmt_float, _fmt_macs"
+    "from fasterbench.core import _bytes_to_mib, _device_ctx, _measuring, _section, _fmt_table, _fmt_float, _fmt_macs"
    ]
   },
   {
@@ -217,7 +217,7 @@
     "    config = _METRIC_CONFIG[metric]\n",
     "    col = config[\"col\"]\n",
     "    \n",
-    "    with _device_ctx(device) as dev, _on_device(model.eval(), dev) as model:\n",
+    "    with _device_ctx(device) as dev, _measuring(model, dev) as model:\n",
     "        sample = sample.to(dev)\n",
     "        \n",
     "        leaf_mods = _leaf_modules(model)\n",
@@ -482,7 +482,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))\n_x = torch.randn(1, 10)\n_lp = LayerProfiler(_m, _x)\n_results = _lp.profile(metrics=[\"size\"])\nassert len(_results) > 0\n\n# profiling leaves the model on the device it came in on\ntest_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})"
+    "#| hide\n",
+    "from fastcore.test import *\n",
+    "\n",
+    "import torch, torch.nn as nn\n",
+    "_m = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5))\n",
+    "_x = torch.randn(1, 10)\n",
+    "_lp = LayerProfiler(_m, _x)\n",
+    "_results = _lp.profile(metrics=[\"size\"])\n",
+    "assert len(_results) > 0\n",
+    "\n",
+    "# profiling leaves the model on the device it came in on\n",
+    "test_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})\n",
+    "\n",
+    "# ...and in the mode it came in\n",
+    "_tm = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5)).train()\n",
+    "LayerProfiler(_tm, _x).profile(metrics=[\"speed\", \"memory\"], warmup=1, steps=2)\n",
+    "test_eq(_tm.training, True)\n",
+    "_em = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5)).eval()\n",
+    "LayerProfiler(_em, _x).profile(metrics=[\"speed\", \"memory\"], warmup=1, steps=2)\n",
+    "test_eq(_em.training, False)"
    ]
   },
   {

nbs/analysis/roofline.ipynb

@@ -50,7 +50,7 @@
     "import torch.nn as nn\n",
     "import plotly.graph_objects as go\n",
     "\n",
-    "from fasterbench.core import _device_ctx, _on_device, _sync, _fmt_float, _section\n",
+    "from fasterbench.core import _device_ctx, _measuring, _sync, _fmt_float, _section\n",
     "from fasterbench.profiling import _leaf_modules, _tensor_bytes, _output_bytes"
    ]
   },
@@ -415,7 +415,7 @@
     "            if self.peaks is None:\n",
     "                self.peaks = measure_peaks(dev, dtype=self.sample.dtype)\n",
     "\n",
-    "            with _on_device(self.model.eval(), dev) as model:\n",
+    "            with _measuring(self.model, dev) as model:\n",
     "                sample = self.sample.to(dev)\n",
     "                leaf_mods = _leaf_modules(model)\n",
     "\n",
@@ -747,6 +747,13 @@
     "    assert math.isfinite(_r.utilization_pct)\n",
     "# profiling leaves the model on the device it came in on\n",
     "test_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})\n",
+    "# ...and in the mode it came in\n",
+    "_tm = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 16)).train()\n",
+    "RooflineAnalyzer(_tm, _x, peaks=_peaks).profile(device=\"cpu\", warmup=1, steps=2)\n",
+    "test_eq(_tm.training, True)\n",
+    "_em = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 16)).eval()\n",
+    "RooflineAnalyzer(_em, _x, peaks=_peaks).profile(device=\"cpu\", warmup=1, steps=2)\n",
+    "test_eq(_em.training, False)\n",
     "# summary and plot should run without error\n",
     "_ra.summary(top=5)\n",
     "_fig = _ra.plot()\n",

nbs/core/core.ipynb

@@ -195,6 +195,25 @@
     "    return out"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "44ce3eab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| export\n",
+    "@contextlib.contextmanager\n",
+    "def _measuring(model: torch.nn.Module, dev: str | torch.device):\n",
+    "    \"\"\"Yield `model` on `dev` in eval mode, then restore its device and training mode.\"\"\"\n",
+    "    was_training = model.training\n",
+    "    try:\n",
+    "        with _on_device(model, dev) as m:\n",
+    "            yield m.eval()\n",
+    "    finally:\n",
+    "        model.train(was_training)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -231,7 +250,25 @@
     "test_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})\n",
     "\n",
     "# a model with neither parameters nor buffers does not raise\n",
-    "with _on_device(nn.ReLU(), \"cpu\"): pass"
+    "with _on_device(nn.ReLU(), \"cpu\"): pass\n",
+    "\n",
+    "# _measuring yields the model in eval mode and restores the mode it came in\n",
+    "_tm = nn.Linear(4, 2).train()\n",
+    "with _measuring(_tm, \"cpu\") as _inner:\n",
+    "    assert _inner is _tm\n",
+    "    test_eq(_inner.training, False)\n",
+    "test_eq(_tm.training, True)\n",
+    "\n",
+    "# ...restored even when the body raises\n",
+    "with ExceptionExpected(ValueError):\n",
+    "    with _measuring(_tm, \"cpu\"):\n",
+    "        raise ValueError(\"boom\")\n",
+    "test_eq(_tm.training, True)\n",
+    "\n",
+    "# an eval-mode model stays in eval mode\n",
+    "_m.eval()\n",
+    "with _measuring(_m, \"cpu\"): pass\n",
+    "test_eq(_m.training, False)"
    ]
   },
   {
@@ -254,7 +291,15 @@
     "    with ExceptionExpected(ValueError):\n",
     "        with _on_device(_cm, \"cpu\"):\n",
     "            raise ValueError(\"boom\")\n",
-    "    test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
+    "    test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+    "\n",
+    "    # _measuring restores device *and* mode\n",
+    "    _cm.train()\n",
+    "    with _measuring(_cm, \"cpu\") as _inner:\n",
+    "        test_eq({p.device.type for p in _inner.parameters()}, {\"cpu\"})\n",
+    "        test_eq(_inner.training, False)\n",
+    "    test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+    "    test_eq(_cm.training, True)"
    ]
   },
   {

nbs/metrics/energy.ipynb

@@ -43,7 +43,7 @@
     "\n",
     "import torch\n",
     "\n",
-    "from fasterbench.core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported\n",
+    "from fasterbench.core import _device_ctx, _measuring, _sync, _run_on_devices, _ensure_device_supported\n",
     "\n",
     "try:\n",
     "    from codecarbon import EmissionsTracker, OfflineEmissionsTracker\n",
@@ -135,7 +135,7 @@
     "\n",
     "    with _device_ctx(device) as dev:\n",
     "        _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV\n",
-    "        with _on_device(model.eval(), dev) as model:\n",
+    "        with _measuring(model, dev) as model:\n",
     "            sample = sample.to(dev, non_blocking=True)\n",
     "\n",
     "            for _ in range(warmup):\n",
@@ -224,6 +224,28 @@
     "show_doc(compute_energy_multi)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b5257aa8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#| hide\n",
+    "#| notest\n",
+    "from fastcore.test import *\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "# measuring leaves the model in the mode it came in\n",
+    "_tm = nn.Linear(10, 5).train()\n",
+    "_x = torch.randn(1, 10)\n",
+    "compute_energy(_tm, _x, warmup=1, steps=2)\n",
+    "test_eq(_tm.training, True)\n",
+    "_em = nn.Linear(10, 5).eval()\n",
+    "compute_energy(_em, _x, warmup=1, steps=2)\n",
+    "test_eq(_em.training, False)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -233,7 +255,7 @@
    "source": [
     "#| hide\n",
     "#| notest\n",
-    "# Runs codecarbon (and needs a second device): measuring leaves the model where it was.\n",
+    "# Runs codecarbon: measuring leaves the model where it was (only the CUDA block needs a second device).\n",
     "from fastcore.test import *\n",
     "import torch.nn as nn\n",
     "\n",