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

fasterbench/__init__.py

@@ -1,16 +1,7 @@
-"""Comprehensive benchmarking toolkit for deep learning models
+"""Comprehensive benchmarking toolkit for deep learning models"""
 
-Docs: https://FasterAI-Labs.github.io/fasterbench/index.html.md"""
-
-# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/index.ipynb.
-
-# %% auto #0
-__all__ = []
-
-# %% ../nbs/index.ipynb #version_cell
 __version__ = "0.1.0"
 
-# %% ../nbs/index.ipynb #8b6f8c52
 from .benchmark import benchmark, BenchmarkResult
 from .size import SizeMetrics, compute_size, get_model_size, get_num_parameters
 from fasterbench.speed import (

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._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.core._sync': ('core/core.html#_sync', 'fasterbench/core.py')},

fasterbench/core.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import contextlib, warnings
+from itertools import chain
 import torch
 
 # %% auto #0
@@ -52,6 +53,19 @@ def _device_ctx(dev: str | torch.device):  # device string or torch.device
         dev = torch.device("cpu")
     yield dev
 
+
+@contextlib.contextmanager
+def _on_device(model: torch.nn.Module, dev: str | torch.device):
+    """Yield `model` on `dev`, then move it back where it was."""
+    t = next(chain(model.parameters(), model.buffers()), None)
+    orig = t.device if t is not None else dev  # no tensors: nothing to restore
+    try:
+        model.to(dev)
+        yield model
+    finally:
+        model.to(orig)
+
+
 def _sync(dev: torch.device) -> None:  # device to synchronize
     """Synchronize CUDA device if applicable."""
     if dev.type == "cuda":

fasterbench/energy.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from .core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported
+from .core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported
 
 try:
     from codecarbon import EmissionsTracker, OfflineEmissionsTracker
@@ -95,22 +95,22 @@ def compute_energy(
 
     with _device_ctx(device) as dev:
         _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-        model = model.eval().to(dev)
-        sample = sample.to(dev, non_blocking=True)
+        with _on_device(model.eval(), dev) as model:
+            sample = sample.to(dev, non_blocking=True)
 
-        for _ in range(warmup):
-            model(sample)
-        _sync(dev)
-
-        tracker.start()
-        try:
-            t0 = time.perf_counter()
-            for _ in range(steps):
+            for _ in range(warmup):
                 model(sample)
             _sync(dev)
-        finally:
-            tracker.stop()
-        dur_s = time.perf_counter() - t0
+
+            tracker.start()
+            try:
+                t0 = time.perf_counter()
+                for _ in range(steps):
+                    model(sample)
+                _sync(dev)
+            finally:
+                tracker.stop()
+            dur_s = time.perf_counter() - t0
 
     # codecarbon silently fails if another instance is running,
     # leaving final_emissions_data as None

fasterbench/memory.py

@@ -6,7 +6,7 @@
 import warnings
 from dataclasses import dataclass, asdict
 from typing import Sequence
-from .core import _bytes_to_mib, _run_on_devices, _ensure_device_supported
+from .core import _bytes_to_mib, _on_device, _run_on_devices, _ensure_device_supported
 
 import numpy as np
 import torch
@@ -50,21 +50,21 @@ def _gpu_metrics(
     """Measure GPU memory usage."""
     dev = torch.device("cuda")
     _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-    model = model.eval().to(dev)
-    sample = sample.to(dev)
+    with _on_device(model.eval(), dev) as model:
+        sample = sample.to(dev)
 
-    for _ in range(warmup):
-        model(sample)
-    torch.cuda.synchronize(dev)
-
-    alloc, alloc_peak, reserv_peak = [], [], []
-    for _ in range(steps):
-        torch.cuda.reset_peak_memory_stats(dev)
-        model(sample)
+        for _ in range(warmup):
+            model(sample)
         torch.cuda.synchronize(dev)
-        alloc.append(torch.cuda.memory_allocated(dev))
-        alloc_peak.append(torch.cuda.max_memory_allocated(dev))
-        reserv_peak.append(torch.cuda.max_memory_reserved(dev))
+
+        alloc, alloc_peak, reserv_peak = [], [], []
+        for _ in range(steps):
+            torch.cuda.reset_peak_memory_stats(dev)
+            model(sample)
+            torch.cuda.synchronize(dev)
+            alloc.append(torch.cuda.memory_allocated(dev))
+            alloc_peak.append(torch.cuda.max_memory_allocated(dev))
+            reserv_peak.append(torch.cuda.max_memory_reserved(dev))
 
     return MemoryMetrics(
         avg_mib=_bytes_to_mib(float(np.mean(alloc))),
@@ -88,22 +88,22 @@ def _cpu_metrics(
         return _nan_memory_metrics("cpu")
 
     proc = psutil.Process()
-    model = model.eval().cpu()
-    sample = sample.cpu()
-
-    rss0 = proc.memory_info().rss
-
-    for _ in range(warmup):
-        model(sample)
-
-    diffs: list[int] = []
-    peaks: list[int] = []
-    for _ in range(steps):
-        rss_before = proc.memory_info().rss
-        model(sample)
-        rss_after = proc.memory_info().rss
-        diffs.append(max(0, rss_after - rss_before))
-        peaks.append(max(0, rss_after - rss0))
+    with _on_device(model.eval(), "cpu") as model:
+        sample = sample.cpu()
+
+        rss0 = proc.memory_info().rss
+
+        for _ in range(warmup):
+            model(sample)
+
+        diffs: list[int] = []
+        peaks: list[int] = []
+        for _ in range(steps):
+            rss_before = proc.memory_info().rss
+            model(sample)
+            rss_after = proc.memory_info().rss
+            diffs.append(max(0, rss_after - rss_before))
+            peaks.append(max(0, rss_after - rss0))
 
     return MemoryMetrics(
         avg_mib=_bytes_to_mib(float(np.mean(diffs))),

fasterbench/profiling.py

@@ -11,7 +11,7 @@
 import torch
 import torch.nn as nn
 
-from .core import _bytes_to_mib, _device_ctx, _section, _fmt_table, _fmt_float, _fmt_macs
+from .core import _bytes_to_mib, _device_ctx, _on_device, _section, _fmt_table, _fmt_float, _fmt_macs
 
 # %% auto #0
 __all__ = ['LayerProfiler']
@@ -177,8 +177,7 @@ def _profile_layers(
     config = _METRIC_CONFIG[metric]
     col = config["col"]
 
-    with _device_ctx(device) as dev:
-        model = model.eval().to(dev)
+    with _device_ctx(device) as dev, _on_device(model.eval(), 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, _sync, _fmt_float, _section
+from .core import _device_ctx, _on_device, _sync, _fmt_float, _section
 from .profiling import _leaf_modules, _tensor_bytes, _output_bytes
 
 # %% auto #0
@@ -304,32 +304,32 @@ def profile(
             if self.peaks is None:
                 self.peaks = measure_peaks(dev, dtype=self.sample.dtype)
 
-            model = self.model.eval().to(dev)
-            sample = self.sample.to(dev)
-            leaf_mods = _leaf_modules(model)
+            with _on_device(self.model.eval(), dev) as model:
+                sample = self.sample.to(dev)
+                leaf_mods = _leaf_modules(model)
 
-            bytes_state: dict[str, list] = {n: [] for n in leaf_mods}
-            time_state: dict[str, list] = {n: [] for n in leaf_mods}
-            flops_state: dict[str, list] = {n: [] for n in leaf_mods}
-            call_state: dict[str, int] = {n: 0 for n in leaf_mods}
+                bytes_state: dict[str, list] = {n: [] for n in leaf_mods}
+                time_state: dict[str, list] = {n: [] for n in leaf_mods}
+                flops_state: dict[str, list] = {n: [] for n in leaf_mods}
+                call_state: dict[str, int] = {n: 0 for n in leaf_mods}
 
-            handles = _setup_roofline_hooks(
-                leaf_mods, bytes_state, time_state, flops_state, call_state, dev.type,
-            )
-            try:
-                for _ in range(warmup):
-                    model(sample)
-                # Reset accumulators after warmup
-                for n in leaf_mods:
-                    bytes_state[n].clear()
-                    time_state[n].clear()
-                    flops_state[n].clear()
-                    call_state[n] = 0
-                for _ in range(steps):
-                    model(sample)
-            finally:
-                for h in handles:
-                    h.remove()
+                handles = _setup_roofline_hooks(
+                    leaf_mods, bytes_state, time_state, flops_state, call_state, dev.type,
+                )
+                try:
+                    for _ in range(warmup):
+                        model(sample)
+                    # Reset accumulators after warmup
+                    for n in leaf_mods:
+                        bytes_state[n].clear()
+                        time_state[n].clear()
+                        flops_state[n].clear()
+                        call_state[n] = 0
+                    for _ in range(steps):
+                        model(sample)
+                finally:
+                    for h in handles:
+                        h.remove()
 
             # Warn if any module was invoked >1x per forward pass (shared module)
             multi_call = [n for n, c in call_state.items() if steps > 0 and c > steps]

fasterbench/speed.py

@@ -11,7 +11,7 @@
 import torch.nn as nn
 from torch.utils.benchmark import Timer
 
-from .core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported
+from .core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported
 
 # %% auto #0
 __all__ = ['SpeedMetrics', 'compute_speed', 'compute_speed_multi', 'sweep_threads', 'sweep_batch_sizes', 'sweep_latency']
@@ -69,33 +69,33 @@ def _forward_latencies(
 
     with _device_ctx(device) as dev:
         _ensure_device_supported(model, dev)  # quantized ops are CPU-only; avoid SIGSEGV
-        model.eval().to(dev)
-        sample = sample.to(dev, non_blocking=True)
-
-        for _ in range(warmup):
-            model(sample)
-        _sync(dev)
-
-        lat: list[float] = []
-        if use_torch_timer and dev.type == "cpu":
-            t = Timer(stmt="model(x)", globals={"model": model, "x": sample})
-            m = t.blocked_autorange(min_run_time=0.3)
-            per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3
-            lat = per_iter.tolist()
-        elif dev.type == "cuda":
-            start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)
-            for _ in range(steps):
-                start_evt.record()
-                model(sample)
-                end_evt.record()
-                _sync(dev)
-                lat.append(start_evt.elapsed_time(end_evt))
-        else:
-            for _ in range(steps):
-                t0 = time.perf_counter()
+        with _on_device(model.eval(), dev) as model:
+            sample = sample.to(dev, non_blocking=True)
+
+            for _ in range(warmup):
                 model(sample)
-                lat.append((time.perf_counter() - t0) * 1e3)
-        return np.asarray(lat, dtype=np.float32)
+            _sync(dev)
+
+            lat: list[float] = []
+            if use_torch_timer and dev.type == "cpu":
+                t = Timer(stmt="model(x)", globals={"model": model, "x": sample})
+                m = t.blocked_autorange(min_run_time=0.3)
+                per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3
+                lat = per_iter.tolist()
+            elif dev.type == "cuda":
+                start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)
+                for _ in range(steps):
+                    start_evt.record()
+                    model(sample)
+                    end_evt.record()
+                    _sync(dev)
+                    lat.append(start_evt.elapsed_time(end_evt))
+            else:
+                for _ in range(steps):
+                    t0 = time.perf_counter()
+                    model(sample)
+                    lat.append((time.perf_counter() - t0) * 1e3)
+            return np.asarray(lat, dtype=np.float32)
 
 
 #| export