diff --git a/fasterbench/__init__.py b/fasterbench/__init__.py
index 560dcee..2a4c305 100644
--- a/fasterbench/__init__.py
+++ b/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 (
diff --git a/fasterbench/_modidx.py b/fasterbench/_modidx.py
index 171e36b..cf6fecf 100644
--- a/fasterbench/_modidx.py
+++ b/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')},
diff --git a/fasterbench/core.py b/fasterbench/core.py
index 7975c16..e4a1541 100644
--- a/fasterbench/core.py
+++ b/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":
diff --git a/fasterbench/energy.py b/fasterbench/energy.py
index 9c5f425..c36732a 100644
--- a/fasterbench/energy.py
+++ b/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
diff --git a/fasterbench/memory.py b/fasterbench/memory.py
index 70fb76f..f18e499 100644
--- a/fasterbench/memory.py
+++ b/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))),
diff --git a/fasterbench/profiling.py b/fasterbench/profiling.py
index 6e87b79..69480fa 100644
--- a/fasterbench/profiling.py
+++ b/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)
diff --git a/fasterbench/roofline.py b/fasterbench/roofline.py
index 8ee82d4..7577d94 100644
--- a/fasterbench/roofline.py
+++ b/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]
diff --git a/fasterbench/speed.py b/fasterbench/speed.py
index 9af3d31..64c3976 100644
--- a/fasterbench/speed.py
+++ b/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
diff --git a/nbs/analysis/benchmark.ipynb b/nbs/analysis/benchmark.ipynb
index a0434eb..ca1e727 100644
--- a/nbs/analysis/benchmark.ipynb
+++ b/nbs/analysis/benchmark.ipynb
@@ -24,7 +24,23 @@
"id": "bbaee268",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nimport json\nimport torch\nfrom collections.abc import Mapping, Iterator\nfrom dataclasses import dataclass, field\nfrom typing import Sequence, Any\n\nfrom fasterbench.core import _fmt_human, _section\nfrom fasterbench.size import *\nfrom fasterbench.speed import *\nfrom fasterbench.compute import *\nfrom fasterbench.memory import *\nfrom fasterbench.energy import *"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "import json\n",
+ "import torch\n",
+ "from collections.abc import Mapping, Iterator\n",
+ "from dataclasses import dataclass, field\n",
+ "from typing import Sequence, Any\n",
+ "\n",
+ "from fasterbench.core import _fmt_human, _section\n",
+ "from fasterbench.size import *\n",
+ "from fasterbench.speed import *\n",
+ "from fasterbench.compute import *\n",
+ "from fasterbench.memory import *\n",
+ "from fasterbench.energy import *"
+ ]
},
{
"cell_type": "code",
@@ -32,7 +48,93 @@
"id": "qwob04wlkqs",
"metadata": {},
"outputs": [],
- "source": "#| export\n@dataclass\nclass BenchmarkResult(Mapping):\n \"\"\"Structured container for benchmark results with typed access.\n \n Provides both typed attribute access (e.g., `result.speed[\"cpu\"].mean_ms`)\n and backward-compatible dict-like access (e.g., `result[\"speed_cpu_mean_ms\"]`).\n \"\"\"\n size: SizeMetrics | None = None\n speed: dict[str, SpeedMetrics] = field(default_factory=dict)\n compute: ComputeMetrics | None = None\n memory: dict[str, MemoryMetrics] = field(default_factory=dict)\n energy: dict[str, EnergyMetrics] = field(default_factory=dict)\n _dict_cache: dict[str, Any] | None = field(default=None, repr=False)\n\n def as_dict(self) -> dict[str, Any]:\n \"\"\"Return flat dict matching the original `benchmark()` return format. Cached.\"\"\"\n if self._dict_cache is not None:\n return self._dict_cache\n \n out: dict[str, Any] = {}\n if self.size is not None:\n out.update({f\"size_{k}\": v for k, v in self.size.as_dict().items()})\n for dev, met in self.speed.items():\n out.update({f\"speed_{dev}_{k}\": v for k, v in met.as_dict().items()})\n if self.compute is not None:\n out.update({f\"compute_{k}\": v for k, v in self.compute.as_dict().items()})\n for dev, met in self.memory.items():\n out.update({f\"memory_{dev}_{k}\": v for k, v in met.as_dict().items()})\n for dev, met in self.energy.items():\n out.update({f\"energy_{dev}_{k}\": v for k, v in met.as_dict().items()})\n \n self._dict_cache = out\n return out\n\n def _format_summary(self) -> str:\n \"\"\"Build the formatted summary string.\"\"\"\n lines = []\n \n if self.size is not None:\n lines.append(_section(\"Size\"))\n lines.append(f\" Disk: {self.size.size_mib:.2f} MiB\")\n lines.append(f\" Params: {_fmt_human(self.size.num_params)}\")\n \n if self.speed:\n lines.append(_section(\"Speed\"))\n for dev, m in self.speed.items():\n lines.append(f\" {dev}: {m.mean_ms:.2f} ms │ {m.throughput_s:.1f} inf/s │ p99: {m.p99_ms:.2f} ms\")\n \n if self.compute is not None and self.compute.macs_available:\n lines.append(_section(\"Compute\"))\n lines.append(f\" MACs: {self.compute.macs_m:.1f} M\")\n \n if self.memory:\n lines.append(_section(\"Memory\"))\n for dev, m in self.memory.items():\n lines.append(f\" {dev}: peak {m.peak_mib:.2f} MiB │ avg {m.avg_mib:.2f} MiB\")\n \n if self.energy:\n lines.append(_section(\"Energy\"))\n for dev, m in self.energy.items():\n lines.append(f\" {dev}: {m.mean_watts:.1f} W │ {m.energy_wh:.4f} Wh/inf │ {m.co2_eq_g:.4f} g CO₂/inf\")\n \n return \"\\n\".join(lines) if lines else \"No metrics measured\"\n\n def summary(self) -> None:\n \"\"\"Print a human-readable summary of benchmark results.\"\"\"\n print(self._format_summary())\n\n def to_dataframe(self):\n \"\"\"Convert results to a pandas DataFrame.\"\"\"\n import pandas as pd\n return pd.DataFrame([self.as_dict()])\n\n def to_json(self) -> str:\n \"\"\"Serialize to JSON string.\"\"\"\n return json.dumps(self.as_dict(), indent=2)\n\n # Mapping protocol - only 3 methods needed, rest inherited\n def __getitem__(self, key: str) -> Any: return self.as_dict()[key]\n def __iter__(self) -> Iterator[str]: return iter(self.as_dict())\n def __len__(self) -> int: return len(self.as_dict())\n\n def __str__(self) -> str: return self._format_summary()\n def __repr__(self) -> str: return self._format_summary()"
+ "source": [
+ "#| export\n",
+ "@dataclass\n",
+ "class BenchmarkResult(Mapping):\n",
+ " \"\"\"Structured container for benchmark results with typed access.\n",
+ " \n",
+ " Provides both typed attribute access (e.g., `result.speed[\"cpu\"].mean_ms`)\n",
+ " and backward-compatible dict-like access (e.g., `result[\"speed_cpu_mean_ms\"]`).\n",
+ " \"\"\"\n",
+ " size: SizeMetrics | None = None\n",
+ " speed: dict[str, SpeedMetrics] = field(default_factory=dict)\n",
+ " compute: ComputeMetrics | None = None\n",
+ " memory: dict[str, MemoryMetrics] = field(default_factory=dict)\n",
+ " energy: dict[str, EnergyMetrics] = field(default_factory=dict)\n",
+ " _dict_cache: dict[str, Any] | None = field(default=None, repr=False)\n",
+ "\n",
+ " def as_dict(self) -> dict[str, Any]:\n",
+ " \"\"\"Return flat dict matching the original `benchmark()` return format. Cached.\"\"\"\n",
+ " if self._dict_cache is not None:\n",
+ " return self._dict_cache\n",
+ " \n",
+ " out: dict[str, Any] = {}\n",
+ " if self.size is not None:\n",
+ " out.update({f\"size_{k}\": v for k, v in self.size.as_dict().items()})\n",
+ " for dev, met in self.speed.items():\n",
+ " out.update({f\"speed_{dev}_{k}\": v for k, v in met.as_dict().items()})\n",
+ " if self.compute is not None:\n",
+ " out.update({f\"compute_{k}\": v for k, v in self.compute.as_dict().items()})\n",
+ " for dev, met in self.memory.items():\n",
+ " out.update({f\"memory_{dev}_{k}\": v for k, v in met.as_dict().items()})\n",
+ " for dev, met in self.energy.items():\n",
+ " out.update({f\"energy_{dev}_{k}\": v for k, v in met.as_dict().items()})\n",
+ " \n",
+ " self._dict_cache = out\n",
+ " return out\n",
+ "\n",
+ " def _format_summary(self) -> str:\n",
+ " \"\"\"Build the formatted summary string.\"\"\"\n",
+ " lines = []\n",
+ " \n",
+ " if self.size is not None:\n",
+ " lines.append(_section(\"Size\"))\n",
+ " lines.append(f\" Disk: {self.size.size_mib:.2f} MiB\")\n",
+ " lines.append(f\" Params: {_fmt_human(self.size.num_params)}\")\n",
+ " \n",
+ " if self.speed:\n",
+ " lines.append(_section(\"Speed\"))\n",
+ " for dev, m in self.speed.items():\n",
+ " lines.append(f\" {dev}: {m.mean_ms:.2f} ms │ {m.throughput_s:.1f} inf/s │ p99: {m.p99_ms:.2f} ms\")\n",
+ " \n",
+ " if self.compute is not None and self.compute.macs_available:\n",
+ " lines.append(_section(\"Compute\"))\n",
+ " lines.append(f\" MACs: {self.compute.macs_m:.1f} M\")\n",
+ " \n",
+ " if self.memory:\n",
+ " lines.append(_section(\"Memory\"))\n",
+ " for dev, m in self.memory.items():\n",
+ " lines.append(f\" {dev}: peak {m.peak_mib:.2f} MiB │ avg {m.avg_mib:.2f} MiB\")\n",
+ " \n",
+ " if self.energy:\n",
+ " lines.append(_section(\"Energy\"))\n",
+ " for dev, m in self.energy.items():\n",
+ " lines.append(f\" {dev}: {m.mean_watts:.1f} W │ {m.energy_wh:.4f} Wh/inf │ {m.co2_eq_g:.4f} g CO₂/inf\")\n",
+ " \n",
+ " return \"\\n\".join(lines) if lines else \"No metrics measured\"\n",
+ "\n",
+ " def summary(self) -> None:\n",
+ " \"\"\"Print a human-readable summary of benchmark results.\"\"\"\n",
+ " print(self._format_summary())\n",
+ "\n",
+ " def to_dataframe(self):\n",
+ " \"\"\"Convert results to a pandas DataFrame.\"\"\"\n",
+ " import pandas as pd\n",
+ " return pd.DataFrame([self.as_dict()])\n",
+ "\n",
+ " def to_json(self) -> str:\n",
+ " \"\"\"Serialize to JSON string.\"\"\"\n",
+ " return json.dumps(self.as_dict(), indent=2)\n",
+ "\n",
+ " # Mapping protocol - only 3 methods needed, rest inherited\n",
+ " def __getitem__(self, key: str) -> Any: return self.as_dict()[key]\n",
+ " def __iter__(self) -> Iterator[str]: return iter(self.as_dict())\n",
+ " def __len__(self) -> int: return len(self.as_dict())\n",
+ "\n",
+ " def __str__(self) -> str: return self._format_summary()\n",
+ " def __repr__(self) -> str: return self._format_summary()"
+ ]
},
{
"cell_type": "code",
@@ -51,7 +153,43 @@
"id": "a89d8062",
"metadata": {},
"outputs": [],
- "source": "#| export\n_VALID_METRICS = frozenset({\"size\", \"speed\", \"compute\", \"memory\", \"energy\"})\n\n\ndef benchmark(\n model: torch.nn.Module, # model to profile\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n metrics: Sequence[str] = (\"size\", \"speed\", \"compute\", \"memory\", \"energy\"), # metrics to compute\n speed_devices: Sequence[str | torch.device] | None = None, # devices for speed (default: cpu + cuda)\n memory_devices: Sequence[str | torch.device] | None = None, # devices for memory\n energy_devices: Sequence[str | torch.device] | None = None, # devices for energy\n **kwargs,\n) -> BenchmarkResult:\n \"\"\"Run comprehensive benchmarks on a model.\"\"\"\n metrics_set = set(metrics)\n if invalid := metrics_set - _VALID_METRICS:\n raise ValueError(f\"Invalid metrics: {invalid}. Valid: {_VALID_METRICS}\")\n if sample.dim() == 0:\n raise ValueError(\"sample must have at least 1 dimension\")\n\n was_training = model.training\n model.eval()\n\n try:\n return BenchmarkResult(\n size=compute_size(model) if \"size\" in metrics_set else None,\n speed=compute_speed_multi(model, sample, devices=speed_devices, **kwargs) if \"speed\" in metrics_set else {},\n compute=compute_compute(model, sample) if \"compute\" in metrics_set else None,\n memory=compute_memory_multi(model, sample, devices=memory_devices, **kwargs) if \"memory\" in metrics_set else {},\n 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()"
+ "source": [
+ "#| export\n",
+ "_VALID_METRICS = frozenset({\"size\", \"speed\", \"compute\", \"memory\", \"energy\"})\n",
+ "\n",
+ "\n",
+ "def benchmark(\n",
+ " model: torch.nn.Module, # model to profile\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " metrics: Sequence[str] = (\"size\", \"speed\", \"compute\", \"memory\", \"energy\"), # metrics to compute\n",
+ " speed_devices: Sequence[str | torch.device] | None = None, # devices for speed (default: cpu + cuda)\n",
+ " memory_devices: Sequence[str | torch.device] | None = None, # devices for memory\n",
+ " energy_devices: Sequence[str | torch.device] | None = None, # devices for energy\n",
+ " **kwargs,\n",
+ ") -> BenchmarkResult:\n",
+ " \"\"\"Run comprehensive benchmarks on a model.\"\"\"\n",
+ " metrics_set = set(metrics)\n",
+ " if invalid := metrics_set - _VALID_METRICS:\n",
+ " raise ValueError(f\"Invalid metrics: {invalid}. Valid: {_VALID_METRICS}\")\n",
+ " if sample.dim() == 0:\n",
+ " raise ValueError(\"sample must have at least 1 dimension\")\n",
+ "\n",
+ " was_training = model.training\n",
+ " model.eval()\n",
+ "\n",
+ " try:\n",
+ " return BenchmarkResult(\n",
+ " size=compute_size(model) if \"size\" in metrics_set else None,\n",
+ " speed=compute_speed_multi(model, sample, devices=speed_devices, **kwargs) if \"speed\" in metrics_set else {},\n",
+ " compute=compute_compute(model, sample) if \"compute\" in metrics_set else None,\n",
+ " memory=compute_memory_multi(model, sample, devices=memory_devices, **kwargs) if \"memory\" in metrics_set else {},\n",
+ " 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()"
+ ]
},
{
"cell_type": "code",
@@ -149,14 +287,70 @@
"id": "075b7b5f",
"metadata": {},
"outputs": [],
- "source": "#| hide\n# Regression: quantized models are CPU-only. PyTorch's quantized ops (fbgemm/qnnpack)\n# have no CUDA kernels, so dispatching them on a CUDA tensor *segfaults* the process.\n# The forward-executing profilers (speed/memory/energy) default to cpu+cuda, so\n# benchmarking a quantized model on a CUDA host must NOT crash — the CUDA leg is\n# skipped (NaN), the CPU leg is real. See: fasterbench/core._is_quantized.\nimport math\nimport torch.ao.quantization as tq\nfrom fasterbench.core import _is_quantized, _ensure_device_supported\n\nclass _QNet(nn.Module):\n \"Tiny static-quantizable net → QuantizedConvReLU2d + QuantizedLinear (pure torch).\"\n def __init__(self):\n super().__init__()\n self.quant, self.dequant = tq.QuantStub(), tq.DeQuantStub()\n self.conv, self.relu = nn.Conv2d(3, 4, 3, padding=1), nn.ReLU()\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(4, 2)\n def forward(self, x):\n x = self.quant(x); x = self.relu(self.conv(x))\n x = self.fc(self.flat(self.pool(x))); return self.dequant(x)\n\n_qm = _QNet().eval()\n_qm.qconfig = tq.get_default_qconfig(\"fbgemm\")\ntq.fuse_modules(_qm, [[\"conv\", \"relu\"]], inplace=True)\ntq.prepare(_qm, inplace=True)\n_qm(torch.randn(4, 3, 8, 8)) # calibrate\ntq.convert(_qm, inplace=True) # -> quantized (CPU-only) modules\nassert _is_quantized(_qm)\n\n_qx = torch.randn(1, 3, 8, 8)\n# Must complete WITHOUT SIGSEGV, even on a CUDA machine.\n_qr = benchmark(_qm, _qx, metrics=[\"size\", \"speed\", \"compute\", \"memory\"], warmup=2, steps=3)\nassert isinstance(_qr, BenchmarkResult)\nassert math.isfinite(_qr.speed[\"cpu\"].mean_ms) # CPU result is real\nassert all(math.isnan(m.mean_ms) for d, m in _qr.speed.items() if d != \"cpu\") # CUDA skipped -> NaN\n\n# Low-level guard: a direct quantized-on-CUDA call raises (catchable), never segfaults.\nwith ExceptionExpected(RuntimeError):\n _ensure_device_supported(_qm, \"cuda\")"
+ "source": [
+ "#| hide\n",
+ "# Regression: quantized models are CPU-only. PyTorch's quantized ops (fbgemm/qnnpack)\n",
+ "# have no CUDA kernels, so dispatching them on a CUDA tensor *segfaults* the process.\n",
+ "# The forward-executing profilers (speed/memory/energy) default to cpu+cuda, so\n",
+ "# benchmarking a quantized model on a CUDA host must NOT crash — the CUDA leg is\n",
+ "# skipped (NaN), the CPU leg is real. See: fasterbench/core._is_quantized.\n",
+ "import math\n",
+ "import torch.ao.quantization as tq\n",
+ "from fasterbench.core import _is_quantized, _ensure_device_supported\n",
+ "\n",
+ "class _QNet(nn.Module):\n",
+ " \"Tiny static-quantizable net → QuantizedConvReLU2d + QuantizedLinear (pure torch).\"\n",
+ " def __init__(self):\n",
+ " super().__init__()\n",
+ " self.quant, self.dequant = tq.QuantStub(), tq.DeQuantStub()\n",
+ " self.conv, self.relu = nn.Conv2d(3, 4, 3, padding=1), nn.ReLU()\n",
+ " self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(4, 2)\n",
+ " def forward(self, x):\n",
+ " x = self.quant(x); x = self.relu(self.conv(x))\n",
+ " x = self.fc(self.flat(self.pool(x))); return self.dequant(x)\n",
+ "\n",
+ "_qm = _QNet().eval()\n",
+ "_qm.qconfig = tq.get_default_qconfig(\"fbgemm\")\n",
+ "tq.fuse_modules(_qm, [[\"conv\", \"relu\"]], inplace=True)\n",
+ "tq.prepare(_qm, inplace=True)\n",
+ "_qm(torch.randn(4, 3, 8, 8)) # calibrate\n",
+ "tq.convert(_qm, inplace=True) # -> quantized (CPU-only) modules\n",
+ "assert _is_quantized(_qm)\n",
+ "\n",
+ "_qx = torch.randn(1, 3, 8, 8)\n",
+ "# Must complete WITHOUT SIGSEGV, even on a CUDA machine.\n",
+ "_qr = benchmark(_qm, _qx, metrics=[\"size\", \"speed\", \"compute\", \"memory\"], warmup=2, steps=3)\n",
+ "assert isinstance(_qr, BenchmarkResult)\n",
+ "assert math.isfinite(_qr.speed[\"cpu\"].mean_ms) # CPU result is real\n",
+ "assert all(math.isnan(m.mean_ms) for d, m in _qr.speed.items() if d != \"cpu\") # CUDA skipped -> NaN\n",
+ "\n",
+ "# Low-level guard: a direct quantized-on-CUDA call raises (catchable), never segfaults.\n",
+ "with ExceptionExpected(RuntimeError):\n",
+ " _ensure_device_supported(_qm, \"cuda\")"
+ ]
},
{
- "cell_type": "markdown",
+ "cell_type": "code",
"execution_count": null,
- "id": "35250fb9",
+ "id": "0997437b",
"metadata": {},
"outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: benchmarking a CUDA model on the CPU must give it back on CUDA,\n",
+ "# otherwise every later timing of that model is silently a CPU timing.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Linear(10, 5).cuda()\n",
+ " benchmark(_cm, _x, metrics=[\"size\", \"speed\", \"compute\", \"memory\"],\n",
+ " speed_devices=[\"cpu\"], memory_devices=[\"cpu\"], warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "35250fb9",
+ "metadata": {},
"source": [
"---\n",
"\n",
diff --git a/nbs/analysis/profiling.ipynb b/nbs/analysis/profiling.ipynb
index ca47f46..caa9f95 100644
--- a/nbs/analysis/profiling.ipynb
+++ b/nbs/analysis/profiling.ipynb
@@ -34,7 +34,20 @@
"id": "imports",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nimport time\nimport warnings\nfrom typing import Sequence\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\n\nfrom fasterbench.core import _bytes_to_mib, _device_ctx, _section, _fmt_table, _fmt_float, _fmt_macs"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "import time\n",
+ "import warnings\n",
+ "from typing import Sequence\n",
+ "\n",
+ "import numpy as np\n",
+ "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"
+ ]
},
{
"cell_type": "code",
@@ -204,8 +217,7 @@
" config = _METRIC_CONFIG[metric]\n",
" col = config[\"col\"]\n",
" \n",
- " with _device_ctx(device) as dev:\n",
- " model = model.eval().to(dev)\n",
+ " with _device_ctx(device) as dev, _on_device(model.eval(), dev) as model:\n",
" sample = sample.to(dev)\n",
" \n",
" leaf_mods = _leaf_modules(model)\n",
@@ -271,7 +283,157 @@
"id": "layer_profiler",
"metadata": {},
"outputs": [],
- "source": "#| export\nclass LayerProfiler:\n \"\"\"Unified per-layer profiler for multiple metrics (speed, memory, size, compute).\"\"\"\n \n VALID_METRICS = frozenset({\"speed\", \"memory\", \"size\", \"compute\"})\n _CONFIG = {\n \"speed\": {\n \"col\": \"speed_ms\", \"pct\": \"speed_percent\", \"unit\": \"ms\",\n \"label\": \"Speed (slowest)\",\n \"src_col\": \"time_ms\",\n \"format\": _fmt_float,\n },\n \"memory\": {\n \"col\": \"memory_mib\", \"pct\": \"memory_percent\", \"unit\": \"MiB\",\n \"label\": \"Memory (largest)\",\n \"src_col\": \"memory_mib\",\n \"format\": _fmt_float,\n },\n \"size\": {\n \"col\": \"params\", \"pct\": \"params_percent\", \"unit\": \"\",\n \"label\": \"Parameters (largest)\",\n \"src_col\": \"params\",\n \"format\": _fmt_table,\n },\n \"compute\": {\n \"col\": \"macs\", \"pct\": \"macs_percent\", \"unit\": \"\",\n \"label\": \"MACs (most)\",\n \"src_col\": \"macs\",\n \"format\": _fmt_macs,\n },\n }\n \n def __init__(\n self,\n model: nn.Module, # model to profile\n sample: torch.Tensor, # input tensor (with batch dimension)\n ):\n self.model = model\n self.sample = sample\n self._leaf_modules = _leaf_modules(model)\n self._results: list[dict] = []\n self._profiled_metrics: list[str] = []\n \n def profile(\n self,\n metrics: str | Sequence[str] = \"speed\", # metrics to profile: speed, memory, size, compute\n *,\n device: str | torch.device = \"cpu\", # device for speed/memory profiling\n warmup: int = 5, # warmup iterations\n steps: int = 20, # measurement iterations\n ) -> list[dict]:\n \"\"\"Profile specified metrics for each layer, returns list of dicts.\"\"\"\n if isinstance(metrics, str):\n metrics = [metrics]\n metrics = list(metrics)\n \n invalid = set(metrics) - self.VALID_METRICS\n if invalid:\n raise ValueError(f\"Invalid metrics: {invalid}. Valid: {self.VALID_METRICS}\")\n \n # Initialize results dict for each leaf module\n results: dict[str, dict] = {\n name: {\"name\": name, \"type\": mod.__class__.__name__}\n for name, mod in self._leaf_modules.items()\n }\n \n # Profile each metric using unified _profile_layers\n for metric in metrics:\n cfg = self._CONFIG[metric]\n profile_data = _profile_layers(\n self.model, self.sample,\n device=device, metric=metric, warmup=warmup, steps=steps\n )\n \n # Merge results\n for r in profile_data:\n name = r[\"name\"]\n if name in results:\n results[name][cfg[\"col\"]] = r[cfg[\"src_col\"]]\n results[name][cfg[\"pct\"]] = r[\"percent\"]\n \n out = list(results.values())\n \n # Sort by first metric\n sort_col = self._CONFIG[metrics[0]][\"col\"]\n out.sort(key=lambda x: x.get(sort_col, 0) or 0, reverse=True)\n \n # Store for top() and summary()\n self._results = out\n self._profiled_metrics = metrics\n \n return out\n \n def top(\n self,\n metric: str, # metric to sort by: speed, memory, size, compute\n n: int = 5, # number of layers to return\n *,\n ascending: bool = False, # if True, return smallest/fastest instead of largest/slowest\n ) -> list[dict]:\n \"\"\"Get top N layers sorted by the specified metric.\"\"\"\n if not self._results:\n raise RuntimeError(\"No results available. Call profile() first.\")\n \n if metric not in self._CONFIG:\n raise ValueError(f\"Invalid metric: {metric}. Valid: {list(self._CONFIG.keys())}\")\n \n col = self._CONFIG[metric][\"col\"]\n \n # Check if metric was profiled\n if col not in self._results[0]:\n raise ValueError(f\"Metric '{metric}' was not profiled. Profiled: {self._profiled_metrics}\")\n \n sorted_results = sorted(\n self._results,\n key=lambda x: x.get(col, 0) or 0,\n reverse=not ascending\n )\n return sorted_results[:n]\n \n def summary(self, *, top: int = 5) -> None:\n \"\"\"Print a formatted summary of top layers for each profiled metric.\"\"\"\n if not self._results:\n raise RuntimeError(\"No results available. Call profile() first.\")\n \n for metric in self._profiled_metrics:\n cfg = self._CONFIG[metric]\n col, pct_col = cfg[\"col\"], cfg[\"pct\"]\n \n # Check if metric data exists\n if col not in self._results[0]:\n continue\n \n print(_section(cfg['label'], 54))\n \n sorted_layers = sorted(\n self._results,\n key=lambda x: x.get(col, 0) or 0,\n reverse=True\n )[:top]\n \n for r in sorted_layers:\n val = r.get(col, 0) or 0\n pct = r.get(pct_col, 0) or 0\n val_str = cfg[\"format\"](val)\n if cfg[\"unit\"]:\n val_str += f\" {cfg['unit']}\"\n print(f\" {r['name']:40} {r['type']:15} {val_str} ({pct:5.1f}%)\")\n print()"
+ "source": [
+ "#| export\n",
+ "class LayerProfiler:\n",
+ " \"\"\"Unified per-layer profiler for multiple metrics (speed, memory, size, compute).\"\"\"\n",
+ " \n",
+ " VALID_METRICS = frozenset({\"speed\", \"memory\", \"size\", \"compute\"})\n",
+ " _CONFIG = {\n",
+ " \"speed\": {\n",
+ " \"col\": \"speed_ms\", \"pct\": \"speed_percent\", \"unit\": \"ms\",\n",
+ " \"label\": \"Speed (slowest)\",\n",
+ " \"src_col\": \"time_ms\",\n",
+ " \"format\": _fmt_float,\n",
+ " },\n",
+ " \"memory\": {\n",
+ " \"col\": \"memory_mib\", \"pct\": \"memory_percent\", \"unit\": \"MiB\",\n",
+ " \"label\": \"Memory (largest)\",\n",
+ " \"src_col\": \"memory_mib\",\n",
+ " \"format\": _fmt_float,\n",
+ " },\n",
+ " \"size\": {\n",
+ " \"col\": \"params\", \"pct\": \"params_percent\", \"unit\": \"\",\n",
+ " \"label\": \"Parameters (largest)\",\n",
+ " \"src_col\": \"params\",\n",
+ " \"format\": _fmt_table,\n",
+ " },\n",
+ " \"compute\": {\n",
+ " \"col\": \"macs\", \"pct\": \"macs_percent\", \"unit\": \"\",\n",
+ " \"label\": \"MACs (most)\",\n",
+ " \"src_col\": \"macs\",\n",
+ " \"format\": _fmt_macs,\n",
+ " },\n",
+ " }\n",
+ " \n",
+ " def __init__(\n",
+ " self,\n",
+ " model: nn.Module, # model to profile\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " ):\n",
+ " self.model = model\n",
+ " self.sample = sample\n",
+ " self._leaf_modules = _leaf_modules(model)\n",
+ " self._results: list[dict] = []\n",
+ " self._profiled_metrics: list[str] = []\n",
+ " \n",
+ " def profile(\n",
+ " self,\n",
+ " metrics: str | Sequence[str] = \"speed\", # metrics to profile: speed, memory, size, compute\n",
+ " *,\n",
+ " device: str | torch.device = \"cpu\", # device for speed/memory profiling\n",
+ " warmup: int = 5, # warmup iterations\n",
+ " steps: int = 20, # measurement iterations\n",
+ " ) -> list[dict]:\n",
+ " \"\"\"Profile specified metrics for each layer, returns list of dicts.\"\"\"\n",
+ " if isinstance(metrics, str):\n",
+ " metrics = [metrics]\n",
+ " metrics = list(metrics)\n",
+ " \n",
+ " invalid = set(metrics) - self.VALID_METRICS\n",
+ " if invalid:\n",
+ " raise ValueError(f\"Invalid metrics: {invalid}. Valid: {self.VALID_METRICS}\")\n",
+ " \n",
+ " # Initialize results dict for each leaf module\n",
+ " results: dict[str, dict] = {\n",
+ " name: {\"name\": name, \"type\": mod.__class__.__name__}\n",
+ " for name, mod in self._leaf_modules.items()\n",
+ " }\n",
+ " \n",
+ " # Profile each metric using unified _profile_layers\n",
+ " for metric in metrics:\n",
+ " cfg = self._CONFIG[metric]\n",
+ " profile_data = _profile_layers(\n",
+ " self.model, self.sample,\n",
+ " device=device, metric=metric, warmup=warmup, steps=steps\n",
+ " )\n",
+ " \n",
+ " # Merge results\n",
+ " for r in profile_data:\n",
+ " name = r[\"name\"]\n",
+ " if name in results:\n",
+ " results[name][cfg[\"col\"]] = r[cfg[\"src_col\"]]\n",
+ " results[name][cfg[\"pct\"]] = r[\"percent\"]\n",
+ " \n",
+ " out = list(results.values())\n",
+ " \n",
+ " # Sort by first metric\n",
+ " sort_col = self._CONFIG[metrics[0]][\"col\"]\n",
+ " out.sort(key=lambda x: x.get(sort_col, 0) or 0, reverse=True)\n",
+ " \n",
+ " # Store for top() and summary()\n",
+ " self._results = out\n",
+ " self._profiled_metrics = metrics\n",
+ " \n",
+ " return out\n",
+ " \n",
+ " def top(\n",
+ " self,\n",
+ " metric: str, # metric to sort by: speed, memory, size, compute\n",
+ " n: int = 5, # number of layers to return\n",
+ " *,\n",
+ " ascending: bool = False, # if True, return smallest/fastest instead of largest/slowest\n",
+ " ) -> list[dict]:\n",
+ " \"\"\"Get top N layers sorted by the specified metric.\"\"\"\n",
+ " if not self._results:\n",
+ " raise RuntimeError(\"No results available. Call profile() first.\")\n",
+ " \n",
+ " if metric not in self._CONFIG:\n",
+ " raise ValueError(f\"Invalid metric: {metric}. Valid: {list(self._CONFIG.keys())}\")\n",
+ " \n",
+ " col = self._CONFIG[metric][\"col\"]\n",
+ " \n",
+ " # Check if metric was profiled\n",
+ " if col not in self._results[0]:\n",
+ " raise ValueError(f\"Metric '{metric}' was not profiled. Profiled: {self._profiled_metrics}\")\n",
+ " \n",
+ " sorted_results = sorted(\n",
+ " self._results,\n",
+ " key=lambda x: x.get(col, 0) or 0,\n",
+ " reverse=not ascending\n",
+ " )\n",
+ " return sorted_results[:n]\n",
+ " \n",
+ " def summary(self, *, top: int = 5) -> None:\n",
+ " \"\"\"Print a formatted summary of top layers for each profiled metric.\"\"\"\n",
+ " if not self._results:\n",
+ " raise RuntimeError(\"No results available. Call profile() first.\")\n",
+ " \n",
+ " for metric in self._profiled_metrics:\n",
+ " cfg = self._CONFIG[metric]\n",
+ " col, pct_col = cfg[\"col\"], cfg[\"pct\"]\n",
+ " \n",
+ " # Check if metric data exists\n",
+ " if col not in self._results[0]:\n",
+ " continue\n",
+ " \n",
+ " print(_section(cfg['label'], 54))\n",
+ " \n",
+ " sorted_layers = sorted(\n",
+ " self._results,\n",
+ " key=lambda x: x.get(col, 0) or 0,\n",
+ " reverse=True\n",
+ " )[:top]\n",
+ " \n",
+ " for r in sorted_layers:\n",
+ " val = r.get(col, 0) or 0\n",
+ " pct = r.get(pct_col, 0) or 0\n",
+ " val_str = cfg[\"format\"](val)\n",
+ " if cfg[\"unit\"]:\n",
+ " val_str += f\" {cfg['unit']}\"\n",
+ " print(f\" {r['name']:40} {r['type']:15} {val_str} ({pct:5.1f}%)\")\n",
+ " print()"
+ ]
},
{
"cell_type": "code",
@@ -320,7 +482,25 @@
"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"
+ "#| 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\"})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3e090fbf",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: profiling a CUDA model on the CPU must not park it there.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Sequential(nn.Linear(10, 20), nn.ReLU(), nn.Linear(20, 5)).cuda()\n",
+ " _profile_layers(_cm, _x, device=\"cpu\", metric=\"speed\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+ " LayerProfiler(_cm, _x).profile(metrics=[\"speed\", \"memory\"], device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
]
},
{
diff --git a/nbs/analysis/roofline.ipynb b/nbs/analysis/roofline.ipynb
index b4d4024..85b334a 100644
--- a/nbs/analysis/roofline.ipynb
+++ b/nbs/analysis/roofline.ipynb
@@ -35,7 +35,24 @@
"id": "imports",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nimport math\nimport time\nimport warnings\nfrom dataclasses import dataclass, asdict\nfrom contextlib import contextmanager\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport plotly.graph_objects as go\n\nfrom fasterbench.core import _device_ctx, _sync, _fmt_float, _section\nfrom fasterbench.profiling import _leaf_modules, _tensor_bytes, _output_bytes"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "import math\n",
+ "import time\n",
+ "import warnings\n",
+ "from dataclasses import dataclass, asdict\n",
+ "from contextlib import contextmanager\n",
+ "\n",
+ "import numpy as np\n",
+ "import torch\n",
+ "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.profiling import _leaf_modules, _tensor_bytes, _output_bytes"
+ ]
},
{
"cell_type": "code",
@@ -247,7 +264,111 @@
"id": "roofline_hooks",
"metadata": {},
"outputs": [],
- "source": "#| export\ndef _layer_flops(module: nn.Module, inp, output) -> float:\n \"Estimate FLOPs for a single leaf module forward call. Returns 0 for layer types we do not model.\"\n # Covers the layer types where FLOPs actually dominate - conv, linear, matmul-adjacent.\n # Cheap element-wise ops (ReLU, BN, pooling) are intentionally assigned 0 FLOPs; they end up\n # in the \\\"undefined\\\" bucket (memory-bound trivially) and are surfaced via the warning.\n t_in = inp[0] if isinstance(inp, tuple) and len(inp) > 0 and isinstance(inp[0], torch.Tensor) else None\n t_out = output if isinstance(output, torch.Tensor) else None\n\n if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)) and t_out is not None:\n cin = module.in_channels // module.groups\n cout = module.out_channels\n ksize = 1\n for k in (module.kernel_size if isinstance(module.kernel_size, tuple) else (module.kernel_size,)):\n ksize *= k\n spatial = 1\n for s in t_out.shape[2:]:\n spatial *= int(s)\n batch = int(t_out.shape[0]) if t_out.ndim > 0 else 1\n macs = batch * cin * cout * ksize * spatial\n flops = 2.0 * macs\n if module.bias is not None:\n flops += batch * cout * spatial # bias add\n return float(flops)\n\n if isinstance(module, (nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d)) and t_in is not None:\n cin = module.in_channels // module.groups\n cout = module.out_channels\n ksize = 1\n for k in (module.kernel_size if isinstance(module.kernel_size, tuple) else (module.kernel_size,)):\n ksize *= k\n spatial = 1\n for s in t_in.shape[2:]:\n spatial *= int(s)\n batch = int(t_in.shape[0]) if t_in.ndim > 0 else 1\n return float(2.0 * batch * cin * cout * ksize * spatial)\n\n if isinstance(module, nn.Linear) and t_in is not None:\n in_features = module.in_features\n out_features = module.out_features\n # Everything before the last dim is treated as batch.\n batch = 1\n for s in t_in.shape[:-1]:\n batch *= int(s)\n flops = 2.0 * batch * in_features * out_features\n if module.bias is not None:\n flops += batch * out_features\n return float(flops)\n\n return 0.0\n\n\n#| export\ndef _setup_roofline_hooks(\n leaf_modules, # {name: module} for leaf modules\n bytes_state, # {name: []} to accumulate bytes per call\n time_state, # {name: []} to accumulate seconds per call\n flops_state, # {name: []} to accumulate FLOPs per call\n call_state, # {name: int} counter of forward calls\n device_type, # \"cuda\" or \"cpu\"\n):\n \"Register hooks to measure (bytes moved, time, FLOPs, call count) per layer.\"\n handles = []\n is_cuda = device_type == \"cuda\"\n\n def make_hooks(name: str, module: nn.Module):\n w_bytes = sum(_tensor_bytes(p) for p in module.parameters(recurse=False))\n if is_cuda:\n state = {\"start\": None, \"end\": None}\n def pre(mod, inp):\n state[\"start\"] = torch.cuda.Event(enable_timing=True)\n state[\"end\"] = torch.cuda.Event(enable_timing=True)\n state[\"start\"].record()\n in_b = sum(_tensor_bytes(t) for t in inp if isinstance(t, torch.Tensor))\n bytes_state[name].append(in_b + w_bytes)\n def post(mod, inp, output):\n state[\"end\"].record()\n out_b = _output_bytes(output)\n bytes_state[name][-1] += out_b\n torch.cuda.synchronize()\n time_state[name].append(state[\"start\"].elapsed_time(state[\"end\"]) / 1000.0)\n flops_state[name].append(_layer_flops(mod, inp, output))\n call_state[name] += 1\n else:\n state = {\"t0\": 0.0}\n def pre(mod, inp):\n state[\"t0\"] = time.perf_counter()\n in_b = sum(_tensor_bytes(t) for t in inp if isinstance(t, torch.Tensor))\n bytes_state[name].append(in_b + w_bytes)\n def post(mod, inp, output):\n dt = time.perf_counter() - state[\"t0\"]\n out_b = _output_bytes(output)\n bytes_state[name][-1] += out_b\n time_state[name].append(dt)\n flops_state[name].append(_layer_flops(mod, inp, output))\n call_state[name] += 1\n return pre, post\n\n for name, mod in leaf_modules.items():\n pre, post = make_hooks(name, mod)\n handles.append(mod.register_forward_pre_hook(pre))\n handles.append(mod.register_forward_hook(post))\n return handles"
+ "source": [
+ "#| export\n",
+ "def _layer_flops(module: nn.Module, inp, output) -> float:\n",
+ " \"Estimate FLOPs for a single leaf module forward call. Returns 0 for layer types we do not model.\"\n",
+ " # Covers the layer types where FLOPs actually dominate - conv, linear, matmul-adjacent.\n",
+ " # Cheap element-wise ops (ReLU, BN, pooling) are intentionally assigned 0 FLOPs; they end up\n",
+ " # in the \\\"undefined\\\" bucket (memory-bound trivially) and are surfaced via the warning.\n",
+ " t_in = inp[0] if isinstance(inp, tuple) and len(inp) > 0 and isinstance(inp[0], torch.Tensor) else None\n",
+ " t_out = output if isinstance(output, torch.Tensor) else None\n",
+ "\n",
+ " if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)) and t_out is not None:\n",
+ " cin = module.in_channels // module.groups\n",
+ " cout = module.out_channels\n",
+ " ksize = 1\n",
+ " for k in (module.kernel_size if isinstance(module.kernel_size, tuple) else (module.kernel_size,)):\n",
+ " ksize *= k\n",
+ " spatial = 1\n",
+ " for s in t_out.shape[2:]:\n",
+ " spatial *= int(s)\n",
+ " batch = int(t_out.shape[0]) if t_out.ndim > 0 else 1\n",
+ " macs = batch * cin * cout * ksize * spatial\n",
+ " flops = 2.0 * macs\n",
+ " if module.bias is not None:\n",
+ " flops += batch * cout * spatial # bias add\n",
+ " return float(flops)\n",
+ "\n",
+ " if isinstance(module, (nn.ConvTranspose1d, nn.ConvTranspose2d, nn.ConvTranspose3d)) and t_in is not None:\n",
+ " cin = module.in_channels // module.groups\n",
+ " cout = module.out_channels\n",
+ " ksize = 1\n",
+ " for k in (module.kernel_size if isinstance(module.kernel_size, tuple) else (module.kernel_size,)):\n",
+ " ksize *= k\n",
+ " spatial = 1\n",
+ " for s in t_in.shape[2:]:\n",
+ " spatial *= int(s)\n",
+ " batch = int(t_in.shape[0]) if t_in.ndim > 0 else 1\n",
+ " return float(2.0 * batch * cin * cout * ksize * spatial)\n",
+ "\n",
+ " if isinstance(module, nn.Linear) and t_in is not None:\n",
+ " in_features = module.in_features\n",
+ " out_features = module.out_features\n",
+ " # Everything before the last dim is treated as batch.\n",
+ " batch = 1\n",
+ " for s in t_in.shape[:-1]:\n",
+ " batch *= int(s)\n",
+ " flops = 2.0 * batch * in_features * out_features\n",
+ " if module.bias is not None:\n",
+ " flops += batch * out_features\n",
+ " return float(flops)\n",
+ "\n",
+ " return 0.0\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _setup_roofline_hooks(\n",
+ " leaf_modules, # {name: module} for leaf modules\n",
+ " bytes_state, # {name: []} to accumulate bytes per call\n",
+ " time_state, # {name: []} to accumulate seconds per call\n",
+ " flops_state, # {name: []} to accumulate FLOPs per call\n",
+ " call_state, # {name: int} counter of forward calls\n",
+ " device_type, # \"cuda\" or \"cpu\"\n",
+ "):\n",
+ " \"Register hooks to measure (bytes moved, time, FLOPs, call count) per layer.\"\n",
+ " handles = []\n",
+ " is_cuda = device_type == \"cuda\"\n",
+ "\n",
+ " def make_hooks(name: str, module: nn.Module):\n",
+ " w_bytes = sum(_tensor_bytes(p) for p in module.parameters(recurse=False))\n",
+ " if is_cuda:\n",
+ " state = {\"start\": None, \"end\": None}\n",
+ " def pre(mod, inp):\n",
+ " state[\"start\"] = torch.cuda.Event(enable_timing=True)\n",
+ " state[\"end\"] = torch.cuda.Event(enable_timing=True)\n",
+ " state[\"start\"].record()\n",
+ " in_b = sum(_tensor_bytes(t) for t in inp if isinstance(t, torch.Tensor))\n",
+ " bytes_state[name].append(in_b + w_bytes)\n",
+ " def post(mod, inp, output):\n",
+ " state[\"end\"].record()\n",
+ " out_b = _output_bytes(output)\n",
+ " bytes_state[name][-1] += out_b\n",
+ " torch.cuda.synchronize()\n",
+ " time_state[name].append(state[\"start\"].elapsed_time(state[\"end\"]) / 1000.0)\n",
+ " flops_state[name].append(_layer_flops(mod, inp, output))\n",
+ " call_state[name] += 1\n",
+ " else:\n",
+ " state = {\"t0\": 0.0}\n",
+ " def pre(mod, inp):\n",
+ " state[\"t0\"] = time.perf_counter()\n",
+ " in_b = sum(_tensor_bytes(t) for t in inp if isinstance(t, torch.Tensor))\n",
+ " bytes_state[name].append(in_b + w_bytes)\n",
+ " def post(mod, inp, output):\n",
+ " dt = time.perf_counter() - state[\"t0\"]\n",
+ " out_b = _output_bytes(output)\n",
+ " bytes_state[name][-1] += out_b\n",
+ " time_state[name].append(dt)\n",
+ " flops_state[name].append(_layer_flops(mod, inp, output))\n",
+ " call_state[name] += 1\n",
+ " return pre, post\n",
+ "\n",
+ " for name, mod in leaf_modules.items():\n",
+ " pre, post = make_hooks(name, mod)\n",
+ " handles.append(mod.register_forward_pre_hook(pre))\n",
+ " handles.append(mod.register_forward_hook(post))\n",
+ " return handles"
+ ]
},
{
"cell_type": "code",
@@ -255,7 +376,219 @@
"id": "roofline_analyzer",
"metadata": {},
"outputs": [],
- "source": "#| export\nclass RooflineAnalyzer:\n \"Per-layer roofline analysis: measure arithmetic intensity and achieved GFLOPs/s against hardware peaks.\"\n\n def __init__(\n self,\n model: nn.Module, # model to analyze\n sample: torch.Tensor, # input tensor (with batch dimension)\n peaks: HardwarePeaks | None = None, # optional precomputed hardware peaks\n ):\n self.model = model\n self.sample = sample\n self.peaks = peaks\n self._results: list[RooflinePoint] = []\n\n @property\n def results(self) -> list[RooflinePoint]:\n \"Per-layer roofline measurements (populated after profile()).\"\n return self._results\n\n @torch.no_grad()\n def profile(\n self,\n *,\n device: str | torch.device = \"cuda\", # device to profile on\n warmup: int = 5, # warmup iterations\n steps: int = 20, # measurement iterations\n ) -> list[RooflinePoint]:\n \"Run per-layer profiling: collect FLOPs, bytes, and time, then classify against peaks.\"\n with _device_ctx(device) as dev:\n if dev.type == \"cuda\":\n torch.cuda.empty_cache()\n _sync(dev)\n\n # Ensure peaks are available\n if self.peaks is None:\n self.peaks = measure_peaks(dev, dtype=self.sample.dtype)\n\n model = self.model.eval().to(dev)\n sample = self.sample.to(dev)\n leaf_mods = _leaf_modules(model)\n\n bytes_state: dict[str, list] = {n: [] for n in leaf_mods}\n time_state: dict[str, list] = {n: [] for n in leaf_mods}\n flops_state: dict[str, list] = {n: [] for n in leaf_mods}\n call_state: dict[str, int] = {n: 0 for n in leaf_mods}\n\n handles = _setup_roofline_hooks(\n leaf_mods, bytes_state, time_state, flops_state, call_state, dev.type,\n )\n try:\n for _ in range(warmup):\n model(sample)\n # Reset accumulators after warmup\n for n in leaf_mods:\n bytes_state[n].clear()\n time_state[n].clear()\n flops_state[n].clear()\n call_state[n] = 0\n for _ in range(steps):\n model(sample)\n finally:\n for h in handles:\n h.remove()\n\n # Warn if any module was invoked >1x per forward pass (shared module)\n multi_call = [n for n, c in call_state.items() if steps > 0 and c > steps]\n if multi_call:\n warnings.warn(\n f\"{len(multi_call)} module(s) were called more than once per forward \"\n f\"pass; their bytes/time are summed across calls. First: {multi_call[:3]}\"\n )\n\n # --- Build RooflinePoint per layer ---\n peak_flops = self.peaks.peak_flops\n peak_bw = self.peaks.peak_bandwidth\n ridge = self.peaks.ridge_point\n\n results: list[RooflinePoint] = []\n for name, mod in leaf_mods.items():\n b_list = bytes_state[name]\n t_list = time_state[name]\n f_list = flops_state[name]\n bytes_moved = float(np.mean(b_list)) if b_list else 0.0\n time_s = float(np.mean(t_list)) if t_list else 0.0\n flops = float(np.mean(f_list)) if f_list else 0.0\n\n if flops == 0 or bytes_moved == 0 or time_s == 0:\n results.append(RooflinePoint(\n name=name, type=mod.__class__.__name__,\n flops=flops, bytes_moved=bytes_moved, time_s=time_s,\n arithmetic_intensity=0.0, achieved_gflops=0.0,\n bound=\"undefined\", utilization_pct=0.0,\n ))\n continue\n\n ai = flops / bytes_moved\n achieved_gflops = flops / time_s / 1e9\n roof_flops = min(peak_flops, ai * peak_bw)\n roof_gflops = roof_flops / 1e9\n bound = \"memory\" if ai < ridge else \"compute\"\n util = (achieved_gflops / roof_gflops * 100) if roof_gflops > 0 else 0.0\n results.append(RooflinePoint(\n name=name, type=mod.__class__.__name__,\n flops=flops, bytes_moved=bytes_moved, time_s=time_s,\n arithmetic_intensity=ai, achieved_gflops=achieved_gflops,\n bound=bound, utilization_pct=util,\n ))\n\n # Warn if any layer had zero flops/bytes/time -> bound is \"undefined\"\n zero_flops = [r.name for r in results if r.bound == \"undefined\"]\n if zero_flops:\n warnings.warn(\n f\"{len(zero_flops)} layer(s) have undefined roofline (zero FLOPs, bytes, \"\n f\"or time). First: {zero_flops[:3]}\"\n )\n\n self._results = results\n return results\n\n def summary(self, *, top: int = 10) -> None:\n \"Print a table of the slowest layers with their roofline metrics.\"\n if not self._results:\n raise RuntimeError(\"No results available. Call profile() first.\")\n print(_section(\"Roofline\", 72))\n header = f\" {'name':32} {'type':14} {'FLOPs':>10} {'bytes':>10} {'AI':>8} {'GFLOPs/s':>10} {'bound':>9} {'util%':>7}\"\n print(header)\n # Sort by measured time, descending (slowest first)\n sorted_rows = sorted(self._results, key=lambda r: r.time_s, reverse=True)[:top]\n for r in sorted_rows:\n flops_str = f\"{r.flops/1e6:>8.2f}M\" if r.flops >= 1e6 else f\"{r.flops:>10.0f}\"\n bytes_str = f\"{r.bytes_moved/1e6:>8.2f}M\" if r.bytes_moved >= 1e6 else f\"{r.bytes_moved:>10.0f}\"\n ai_str = _fmt_float(r.arithmetic_intensity, width=8, decimals=2)\n gf_str = _fmt_float(r.achieved_gflops, width=10, decimals=2)\n util_str = _fmt_float(r.utilization_pct, width=6, decimals=1) + \"%\"\n print(f\" {r.name:32} {r.type:14} {flops_str} {bytes_str} {ai_str} {gf_str} {r.bound:>9} {util_str}\")\n\n def plot(\n self,\n *,\n title: str = \"Roofline\", # figure title\n ) -> go.Figure:\n \"Render the roofline plot with per-layer scatter points on a log-log grid. Markers are colored by bound classification.\"\n if not self._results:\n raise RuntimeError(\"No results available. Call profile() first.\")\n if self.peaks is None:\n raise RuntimeError(\"No hardware peaks available.\")\n\n peak_flops = self.peaks.peak_flops\n peak_bw = self.peaks.peak_bandwidth\n ridge = self.peaks.ridge_point\n peak_gflops = peak_flops / 1e9\n\n # Build the roof curve: y = min(peak_flops, AI * peak_bw) / 1e9\n valid = [r for r in self._results if r.bound != \"undefined\"]\n if valid:\n ai_min = max(min(r.arithmetic_intensity for r in valid) / 10.0, 1e-3)\n ai_max = max(r.arithmetic_intensity for r in valid) * 10.0\n else:\n ai_min, ai_max = 1e-2, 1e3\n ai_max = max(ai_max, ridge * 10.0)\n\n ai_grid = np.logspace(math.log10(ai_min), math.log10(ai_max), 200)\n roof_gflops = np.minimum(peak_gflops, ai_grid * peak_bw / 1e9)\n\n fig = go.Figure()\n fig.add_trace(go.Scatter(\n x=ai_grid, y=roof_gflops, mode=\"lines\",\n line=dict(color=\"#008080\", width=2),\n name=f\"Roof ({peak_gflops:.0f} GFLOPs/s, {peak_bw/1e9:.1f} GB/s)\",\n hoverinfo=\"skip\",\n ))\n # Ridge point marker\n fig.add_trace(go.Scatter(\n x=[ridge], y=[peak_gflops], mode=\"markers\",\n marker=dict(color=\"#008080\", size=10, symbol=\"diamond\"),\n name=f\"Ridge point ({ridge:.1f} FLOP/byte)\",\n hovertemplate=\"Ridge: %{x:.2f} FLOP/byte\",\n ))\n\n # Color palette for layers\n color_map = {\"memory\": \"#89d6c9\", \"compute\": \"#008080\"}\n for bound_label in (\"memory\", \"compute\"):\n pts = [r for r in valid if r.bound == bound_label]\n if not pts:\n continue\n hover = [\n f\"{r.name}
{r.type}
AI: {r.arithmetic_intensity:.3f} FLOP/byte\"\n f\"
{r.achieved_gflops:.2f} GFLOPs/s
util: {r.utilization_pct:.1f}%\"\n for r in pts\n ]\n fig.add_trace(go.Scatter(\n x=[r.arithmetic_intensity for r in pts],\n y=[r.achieved_gflops for r in pts],\n mode=\"markers\",\n marker=dict(color=color_map[bound_label], size=8, opacity=0.8,\n line=dict(color=\"#008080\", width=0.5)),\n name=f\"{bound_label}-bound\",\n text=hover,\n hovertemplate=\"%{text}\",\n ))\n\n fig.update_layout(\n title=title,\n xaxis=dict(title=\"Arithmetic intensity (FLOP/byte)\", type=\"log\"),\n yaxis=dict(title=\"Achieved performance (GFLOPs/s)\", type=\"log\"),\n paper_bgcolor=\"rgba(0,0,0,0)\",\n plot_bgcolor=\"rgba(0,0,0,0)\",\n legend=dict(bgcolor=\"rgba(0,0,0,0)\"),\n )\n return fig"
+ "source": [
+ "#| export\n",
+ "class RooflineAnalyzer:\n",
+ " \"Per-layer roofline analysis: measure arithmetic intensity and achieved GFLOPs/s against hardware peaks.\"\n",
+ "\n",
+ " def __init__(\n",
+ " self,\n",
+ " model: nn.Module, # model to analyze\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " peaks: HardwarePeaks | None = None, # optional precomputed hardware peaks\n",
+ " ):\n",
+ " self.model = model\n",
+ " self.sample = sample\n",
+ " self.peaks = peaks\n",
+ " self._results: list[RooflinePoint] = []\n",
+ "\n",
+ " @property\n",
+ " def results(self) -> list[RooflinePoint]:\n",
+ " \"Per-layer roofline measurements (populated after profile()).\"\n",
+ " return self._results\n",
+ "\n",
+ " @torch.no_grad()\n",
+ " def profile(\n",
+ " self,\n",
+ " *,\n",
+ " device: str | torch.device = \"cuda\", # device to profile on\n",
+ " warmup: int = 5, # warmup iterations\n",
+ " steps: int = 20, # measurement iterations\n",
+ " ) -> list[RooflinePoint]:\n",
+ " \"Run per-layer profiling: collect FLOPs, bytes, and time, then classify against peaks.\"\n",
+ " with _device_ctx(device) as dev:\n",
+ " if dev.type == \"cuda\":\n",
+ " torch.cuda.empty_cache()\n",
+ " _sync(dev)\n",
+ "\n",
+ " # Ensure peaks are available\n",
+ " 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",
+ " sample = self.sample.to(dev)\n",
+ " leaf_mods = _leaf_modules(model)\n",
+ "\n",
+ " bytes_state: dict[str, list] = {n: [] for n in leaf_mods}\n",
+ " time_state: dict[str, list] = {n: [] for n in leaf_mods}\n",
+ " flops_state: dict[str, list] = {n: [] for n in leaf_mods}\n",
+ " call_state: dict[str, int] = {n: 0 for n in leaf_mods}\n",
+ "\n",
+ " handles = _setup_roofline_hooks(\n",
+ " leaf_mods, bytes_state, time_state, flops_state, call_state, dev.type,\n",
+ " )\n",
+ " try:\n",
+ " for _ in range(warmup):\n",
+ " model(sample)\n",
+ " # Reset accumulators after warmup\n",
+ " for n in leaf_mods:\n",
+ " bytes_state[n].clear()\n",
+ " time_state[n].clear()\n",
+ " flops_state[n].clear()\n",
+ " call_state[n] = 0\n",
+ " for _ in range(steps):\n",
+ " model(sample)\n",
+ " finally:\n",
+ " for h in handles:\n",
+ " h.remove()\n",
+ "\n",
+ " # Warn if any module was invoked >1x per forward pass (shared module)\n",
+ " multi_call = [n for n, c in call_state.items() if steps > 0 and c > steps]\n",
+ " if multi_call:\n",
+ " warnings.warn(\n",
+ " f\"{len(multi_call)} module(s) were called more than once per forward \"\n",
+ " f\"pass; their bytes/time are summed across calls. First: {multi_call[:3]}\"\n",
+ " )\n",
+ "\n",
+ " # --- Build RooflinePoint per layer ---\n",
+ " peak_flops = self.peaks.peak_flops\n",
+ " peak_bw = self.peaks.peak_bandwidth\n",
+ " ridge = self.peaks.ridge_point\n",
+ "\n",
+ " results: list[RooflinePoint] = []\n",
+ " for name, mod in leaf_mods.items():\n",
+ " b_list = bytes_state[name]\n",
+ " t_list = time_state[name]\n",
+ " f_list = flops_state[name]\n",
+ " bytes_moved = float(np.mean(b_list)) if b_list else 0.0\n",
+ " time_s = float(np.mean(t_list)) if t_list else 0.0\n",
+ " flops = float(np.mean(f_list)) if f_list else 0.0\n",
+ "\n",
+ " if flops == 0 or bytes_moved == 0 or time_s == 0:\n",
+ " results.append(RooflinePoint(\n",
+ " name=name, type=mod.__class__.__name__,\n",
+ " flops=flops, bytes_moved=bytes_moved, time_s=time_s,\n",
+ " arithmetic_intensity=0.0, achieved_gflops=0.0,\n",
+ " bound=\"undefined\", utilization_pct=0.0,\n",
+ " ))\n",
+ " continue\n",
+ "\n",
+ " ai = flops / bytes_moved\n",
+ " achieved_gflops = flops / time_s / 1e9\n",
+ " roof_flops = min(peak_flops, ai * peak_bw)\n",
+ " roof_gflops = roof_flops / 1e9\n",
+ " bound = \"memory\" if ai < ridge else \"compute\"\n",
+ " util = (achieved_gflops / roof_gflops * 100) if roof_gflops > 0 else 0.0\n",
+ " results.append(RooflinePoint(\n",
+ " name=name, type=mod.__class__.__name__,\n",
+ " flops=flops, bytes_moved=bytes_moved, time_s=time_s,\n",
+ " arithmetic_intensity=ai, achieved_gflops=achieved_gflops,\n",
+ " bound=bound, utilization_pct=util,\n",
+ " ))\n",
+ "\n",
+ " # Warn if any layer had zero flops/bytes/time -> bound is \"undefined\"\n",
+ " zero_flops = [r.name for r in results if r.bound == \"undefined\"]\n",
+ " if zero_flops:\n",
+ " warnings.warn(\n",
+ " f\"{len(zero_flops)} layer(s) have undefined roofline (zero FLOPs, bytes, \"\n",
+ " f\"or time). First: {zero_flops[:3]}\"\n",
+ " )\n",
+ "\n",
+ " self._results = results\n",
+ " return results\n",
+ "\n",
+ " def summary(self, *, top: int = 10) -> None:\n",
+ " \"Print a table of the slowest layers with their roofline metrics.\"\n",
+ " if not self._results:\n",
+ " raise RuntimeError(\"No results available. Call profile() first.\")\n",
+ " print(_section(\"Roofline\", 72))\n",
+ " header = f\" {'name':32} {'type':14} {'FLOPs':>10} {'bytes':>10} {'AI':>8} {'GFLOPs/s':>10} {'bound':>9} {'util%':>7}\"\n",
+ " print(header)\n",
+ " # Sort by measured time, descending (slowest first)\n",
+ " sorted_rows = sorted(self._results, key=lambda r: r.time_s, reverse=True)[:top]\n",
+ " for r in sorted_rows:\n",
+ " flops_str = f\"{r.flops/1e6:>8.2f}M\" if r.flops >= 1e6 else f\"{r.flops:>10.0f}\"\n",
+ " bytes_str = f\"{r.bytes_moved/1e6:>8.2f}M\" if r.bytes_moved >= 1e6 else f\"{r.bytes_moved:>10.0f}\"\n",
+ " ai_str = _fmt_float(r.arithmetic_intensity, width=8, decimals=2)\n",
+ " gf_str = _fmt_float(r.achieved_gflops, width=10, decimals=2)\n",
+ " util_str = _fmt_float(r.utilization_pct, width=6, decimals=1) + \"%\"\n",
+ " print(f\" {r.name:32} {r.type:14} {flops_str} {bytes_str} {ai_str} {gf_str} {r.bound:>9} {util_str}\")\n",
+ "\n",
+ " def plot(\n",
+ " self,\n",
+ " *,\n",
+ " title: str = \"Roofline\", # figure title\n",
+ " ) -> go.Figure:\n",
+ " \"Render the roofline plot with per-layer scatter points on a log-log grid. Markers are colored by bound classification.\"\n",
+ " if not self._results:\n",
+ " raise RuntimeError(\"No results available. Call profile() first.\")\n",
+ " if self.peaks is None:\n",
+ " raise RuntimeError(\"No hardware peaks available.\")\n",
+ "\n",
+ " peak_flops = self.peaks.peak_flops\n",
+ " peak_bw = self.peaks.peak_bandwidth\n",
+ " ridge = self.peaks.ridge_point\n",
+ " peak_gflops = peak_flops / 1e9\n",
+ "\n",
+ " # Build the roof curve: y = min(peak_flops, AI * peak_bw) / 1e9\n",
+ " valid = [r for r in self._results if r.bound != \"undefined\"]\n",
+ " if valid:\n",
+ " ai_min = max(min(r.arithmetic_intensity for r in valid) / 10.0, 1e-3)\n",
+ " ai_max = max(r.arithmetic_intensity for r in valid) * 10.0\n",
+ " else:\n",
+ " ai_min, ai_max = 1e-2, 1e3\n",
+ " ai_max = max(ai_max, ridge * 10.0)\n",
+ "\n",
+ " ai_grid = np.logspace(math.log10(ai_min), math.log10(ai_max), 200)\n",
+ " roof_gflops = np.minimum(peak_gflops, ai_grid * peak_bw / 1e9)\n",
+ "\n",
+ " fig = go.Figure()\n",
+ " fig.add_trace(go.Scatter(\n",
+ " x=ai_grid, y=roof_gflops, mode=\"lines\",\n",
+ " line=dict(color=\"#008080\", width=2),\n",
+ " name=f\"Roof ({peak_gflops:.0f} GFLOPs/s, {peak_bw/1e9:.1f} GB/s)\",\n",
+ " hoverinfo=\"skip\",\n",
+ " ))\n",
+ " # Ridge point marker\n",
+ " fig.add_trace(go.Scatter(\n",
+ " x=[ridge], y=[peak_gflops], mode=\"markers\",\n",
+ " marker=dict(color=\"#008080\", size=10, symbol=\"diamond\"),\n",
+ " name=f\"Ridge point ({ridge:.1f} FLOP/byte)\",\n",
+ " hovertemplate=\"Ridge: %{x:.2f} FLOP/byte\",\n",
+ " ))\n",
+ "\n",
+ " # Color palette for layers\n",
+ " color_map = {\"memory\": \"#89d6c9\", \"compute\": \"#008080\"}\n",
+ " for bound_label in (\"memory\", \"compute\"):\n",
+ " pts = [r for r in valid if r.bound == bound_label]\n",
+ " if not pts:\n",
+ " continue\n",
+ " hover = [\n",
+ " f\"{r.name}
{r.type}
AI: {r.arithmetic_intensity:.3f} FLOP/byte\"\n",
+ " f\"
{r.achieved_gflops:.2f} GFLOPs/s
util: {r.utilization_pct:.1f}%\"\n",
+ " for r in pts\n",
+ " ]\n",
+ " fig.add_trace(go.Scatter(\n",
+ " x=[r.arithmetic_intensity for r in pts],\n",
+ " y=[r.achieved_gflops for r in pts],\n",
+ " mode=\"markers\",\n",
+ " marker=dict(color=color_map[bound_label], size=8, opacity=0.8,\n",
+ " line=dict(color=\"#008080\", width=0.5)),\n",
+ " name=f\"{bound_label}-bound\",\n",
+ " text=hover,\n",
+ " hovertemplate=\"%{text}\",\n",
+ " ))\n",
+ "\n",
+ " fig.update_layout(\n",
+ " title=title,\n",
+ " xaxis=dict(title=\"Arithmetic intensity (FLOP/byte)\", type=\"log\"),\n",
+ " yaxis=dict(title=\"Achieved performance (GFLOPs/s)\", type=\"log\"),\n",
+ " paper_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " plot_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " legend=dict(bgcolor=\"rgba(0,0,0,0)\"),\n",
+ " )\n",
+ " return fig"
+ ]
},
{
"cell_type": "code",
@@ -314,7 +647,7 @@
"fig.show()\n",
"```\n",
"\n",
- "This is a measurement primitive. Downstream compression workflows (see fasterrecipes) can consume `ra.results` to make decisions - fasterbench itself never prescribes."
+ "This is a measurement primitive: it measures, it never prescribes. The decisions that consume `ra.results` live in the compression workflow that calls it."
]
},
{
@@ -361,7 +694,33 @@
"id": "test_conv_hand_computed",
"metadata": {},
"outputs": [],
- "source": "#| hide\n# Hand-computed Conv2d test: wrap in Sequential so leaf name is non-empty\nimport torch\nimport torch.nn as nn\n\n_synth_peaks = HardwarePeaks(\n peak_flops=1e12, peak_bandwidth=1e11, ridge_point=10.0,\n device=\"cpu\", dtype=\"torch.float32\",\n tf32_enabled=False, cudnn_benchmark=False,\n)\n_conv = nn.Conv2d(4, 8, kernel_size=3, padding=1, bias=False)\n_x = torch.randn(1, 4, 8, 8)\n_model = nn.Sequential(_conv)\n_ra = RooflineAnalyzer(_model, _x, peaks=_synth_peaks)\n_res = _ra.profile(device=\"cpu\", warmup=1, steps=2)\nassert len(_res) == 1\n_r = _res[0]\n# Expected MACs = 4*8*3*3 * 8*8 = 18432; FLOPs = 2*MACs = 36864\ntest_eq(_r.flops, 36864.0)\n# Expected bytes: weights 4*8*3*3*4 = 1152, input 1*4*8*8*4 = 1024, output 1*8*8*8*4 = 2048\n# total = 4224\ntest_eq(_r.bytes_moved, 4224.0)\ntest_close(_r.arithmetic_intensity, 36864.0 / 4224.0, eps=1e-3)\nassert _r.bound in (\"memory\", \"compute\")\nassert math.isfinite(_r.utilization_pct)"
+ "source": [
+ "#| hide\n",
+ "# Hand-computed Conv2d test: wrap in Sequential so leaf name is non-empty\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "_synth_peaks = HardwarePeaks(\n",
+ " peak_flops=1e12, peak_bandwidth=1e11, ridge_point=10.0,\n",
+ " device=\"cpu\", dtype=\"torch.float32\",\n",
+ " tf32_enabled=False, cudnn_benchmark=False,\n",
+ ")\n",
+ "_conv = nn.Conv2d(4, 8, kernel_size=3, padding=1, bias=False)\n",
+ "_x = torch.randn(1, 4, 8, 8)\n",
+ "_model = nn.Sequential(_conv)\n",
+ "_ra = RooflineAnalyzer(_model, _x, peaks=_synth_peaks)\n",
+ "_res = _ra.profile(device=\"cpu\", warmup=1, steps=2)\n",
+ "assert len(_res) == 1\n",
+ "_r = _res[0]\n",
+ "# Expected MACs = 4*8*3*3 * 8*8 = 18432; FLOPs = 2*MACs = 36864\n",
+ "test_eq(_r.flops, 36864.0)\n",
+ "# Expected bytes: weights 4*8*3*3*4 = 1152, input 1*4*8*8*4 = 1024, output 1*8*8*8*4 = 2048\n",
+ "# total = 4224\n",
+ "test_eq(_r.bytes_moved, 4224.0)\n",
+ "test_close(_r.arithmetic_intensity, 36864.0 / 4224.0, eps=1e-3)\n",
+ "assert _r.bound in (\"memory\", \"compute\")\n",
+ "assert math.isfinite(_r.utilization_pct)"
+ ]
},
{
"cell_type": "code",
@@ -386,6 +745,8 @@
" assert _r.arithmetic_intensity >= 0\n",
" assert _r.bound in {\"memory\", \"compute\", \"undefined\"}\n",
" 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",
"# summary and plot should run without error\n",
"_ra.summary(top=5)\n",
"_fig = _ra.plot()\n",
@@ -398,7 +759,26 @@
"id": "test_slow_resnet",
"metadata": {},
"outputs": [],
- "source": "#| hide\n#| notest\nfrom torchvision.models import resnet18\n\n_model = resnet18()\n_sample = torch.randn(1, 3, 64, 64)\n_synth = HardwarePeaks(\n peak_flops=1e13, peak_bandwidth=5e11, ridge_point=20.0,\n device=\"cpu\", dtype=\"torch.float32\",\n tf32_enabled=False, cudnn_benchmark=False,\n)\n_ra = RooflineAnalyzer(_model, _sample, peaks=_synth)\n_results = _ra.profile(device=\"cpu\", warmup=2, steps=3)\nassert len(_results) > 0\nassert all(r.bound in {\"memory\", \"compute\", \"undefined\"} for r in _results)\n_ra.summary(top=5)\n_fig = _ra.plot()\nassert isinstance(_fig, go.Figure)"
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "from torchvision.models import resnet18\n",
+ "\n",
+ "_model = resnet18()\n",
+ "_sample = torch.randn(1, 3, 64, 64)\n",
+ "_synth = HardwarePeaks(\n",
+ " peak_flops=1e13, peak_bandwidth=5e11, ridge_point=20.0,\n",
+ " device=\"cpu\", dtype=\"torch.float32\",\n",
+ " tf32_enabled=False, cudnn_benchmark=False,\n",
+ ")\n",
+ "_ra = RooflineAnalyzer(_model, _sample, peaks=_synth)\n",
+ "_results = _ra.profile(device=\"cpu\", warmup=2, steps=3)\n",
+ "assert len(_results) > 0\n",
+ "assert all(r.bound in {\"memory\", \"compute\", \"undefined\"} for r in _results)\n",
+ "_ra.summary(top=5)\n",
+ "_fig = _ra.plot()\n",
+ "assert isinstance(_fig, go.Figure)"
+ ]
},
{
"cell_type": "code",
@@ -406,13 +786,46 @@
"id": "test_slow_cuda",
"metadata": {},
"outputs": [],
- "source": "#| hide\n#| notest\nif torch.cuda.is_available():\n _p = measure_peaks(device=\"cuda\", matmul_size=512, bandwidth_mb=64, steps=3, warmup=1, cache=False)\n assert _p.peak_flops > 0\n assert _p.peak_bandwidth > 0\n assert _p.ridge_point > 0"
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "if torch.cuda.is_available():\n",
+ " _p = measure_peaks(device=\"cuda\", matmul_size=512, bandwidth_mb=64, steps=3, warmup=1, cache=False)\n",
+ " assert _p.peak_flops > 0\n",
+ " assert _p.peak_bandwidth > 0\n",
+ " assert _p.ridge_point > 0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "b41b65e6",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: profiling a CUDA model on the CPU must not park it there.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 16)).cuda()\n",
+ " _cra = RooflineAnalyzer(_cm, torch.randn(1, 32), peaks=_peaks)\n",
+ " _cra.profile(device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
+ ]
},
{
"cell_type": "markdown",
"id": "see_also",
"metadata": {},
- "source": "---\n\n## See Also\n\n- [Per-layer profiling](profiling.html) - Generic per-layer hook infrastructure reused here\n- [Compute metrics](../metrics/compute.html) - Model-level FLOPs counting\n- [Speed metrics](../metrics/speed.html) - Latency measurement"
+ "source": [
+ "---\n",
+ "\n",
+ "## See Also\n",
+ "\n",
+ "- [Per-layer profiling](profiling.html) - Generic per-layer hook infrastructure reused here\n",
+ "- [Compute metrics](../metrics/compute.html) - Model-level FLOPs counting\n",
+ "- [Speed metrics](../metrics/speed.html) - Latency measurement"
+ ]
}
],
"metadata": {
diff --git a/nbs/core/core.ipynb b/nbs/core/core.ipynb
index c53a5e9..5cebadb 100644
--- a/nbs/core/core.ipynb
+++ b/nbs/core/core.ipynb
@@ -39,6 +39,7 @@
"from __future__ import annotations\n",
"\n",
"import contextlib, warnings\n",
+ "from itertools import chain\n",
"import torch"
]
},
@@ -48,7 +49,151 @@
"id": "ed4d37d2",
"metadata": {},
"outputs": [],
- "source": "#| export\ndef _bytes_to_mib(n: int) -> float: # number of bytes to convert\n \"\"\"Convert bytes to MiB (base-2).\"\"\"\n return n / 1024 ** 2\n\ndef _fmt_human(n: int | float) -> str: # number to format\n \"\"\"Format large number with K/M/B suffix for human-readable output.\"\"\"\n if n >= 1e9: return f\"{n/1e9:.2f}B\"\n if n >= 1e6: return f\"{n/1e6:.2f}M\"\n if n >= 1e3: return f\"{n/1e3:.2f}K\"\n return str(int(n))\n\n\ndef _fmt_table(n: int | float) -> str: # number to format\n \"\"\"Format number with commas for tabular output.\"\"\"\n return f\"{int(n):>12,}\"\n\n\ndef _fmt_float(v: float, width: int = 8, decimals: int = 3) -> str: # value, width, decimal places\n \"\"\"Format float with specified width and decimal places.\"\"\"\n return f\"{v:{width}.{decimals}f}\"\n\n\ndef _fmt_macs(v: float) -> str: # MACs value to format\n \"\"\"Format MACs with appropriate M/K suffix.\"\"\"\n if v >= 1e6: return f\"{v/1e6:8.2f} M\"\n if v >= 1e3: return f\"{v/1e3:8.2f} K\"\n return f\"{int(v):>10}\"\n\n\ndef _section(title: str, width: int = 40) -> str: # section title, total width\n \"\"\"Create a section header with box-drawing characters.\"\"\"\n return f\"═══ {title} \" + \"═\" * (width - len(title))\n\n@contextlib.contextmanager\ndef _device_ctx(dev: str | torch.device): # device string or torch.device\n \"\"\"Context manager that validates device availability and yields resolved device.\"\"\"\n dev = torch.device(dev)\n if dev.type == \"cuda\" and not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n dev = torch.device(\"cpu\")\n yield dev\n\ndef _sync(dev: torch.device) -> None: # device to synchronize\n \"\"\"Synchronize CUDA device if applicable.\"\"\"\n if dev.type == \"cuda\":\n torch.cuda.synchronize(dev)\n\n\ndef _default_devices() -> list[str]:\n \"\"\"Return default device list: ['cpu'] + ['cuda'] if available.\"\"\"\n devices = [\"cpu\"]\n if torch.cuda.is_available():\n devices.append(\"cuda\")\n return devices\n\n\ndef _is_quantized(model: torch.nn.Module) -> bool: # model to inspect\n \"\"\"True if the model contains PyTorch quantized modules.\n\n PyTorch's eager/FX quantized ops (fbgemm/qnnpack) are **CPU-only**. Running\n them on a CUDA tensor has no registered kernel and *segfaults* the process\n rather than raising, so callers must avoid dispatching them off-CPU.\n \"\"\"\n return any(\"quantized\" in type(m).__module__ for m in model.modules())\n\n\ndef _device_supported(\n model: torch.nn.Module, # model to run\n device: str | torch.device, # target device\n) -> bool:\n \"\"\"False if `model` cannot execute on `device` (quantized models are CPU-only).\"\"\"\n return not (torch.device(device).type != \"cpu\" and _is_quantized(model))\n\n\ndef _ensure_device_supported(\n model: torch.nn.Module, # model to run\n device: str | torch.device, # target device\n) -> None:\n \"\"\"Raise a *catchable* error instead of letting PyTorch segfault.\n\n Guards the forward-executing profilers: a quantized model dispatched on a\n non-CPU device would crash the interpreter with SIGSEGV, so we fail loudly\n (and recoverably) first.\n \"\"\"\n if not _device_supported(model, device):\n raise RuntimeError(\n f\"Quantized models run on CPU only; cannot benchmark on '{device}'. \"\n \"PyTorch quantized ops (fbgemm/qnnpack) have no CUDA kernels.\"\n )\n\n\ndef _run_on_devices(\n compute_fn, # single-device compute function\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor\n devices: list[str | torch.device] | None, # devices to run on (None = default)\n nan_factory, # callable(device_str) -> metrics with NaN values\n metric_name: str, # for warning messages (e.g., \"Speed\", \"Memory\")\n **kwargs,\n) -> dict:\n \"\"\"Run compute_fn on multiple devices with unified error handling.\n \n Returns dict mapping device string to metrics (or NaN metrics on failure).\n Devices a model cannot run on (e.g. CUDA for a quantized model) are skipped\n with NaN metrics — this avoids an uncatchable SIGSEGV from dispatching\n CPU-only quantized ops on a CUDA tensor.\n \"\"\"\n if devices is None:\n devices = _default_devices()\n\n out = {}\n for d in devices:\n d_str = str(d)\n if not _device_supported(model, d):\n warnings.warn(\n f\"{metric_name} benchmark skipped on {d}: quantized models run on \"\n \"CPU only (PyTorch quantized ops have no CUDA kernels).\"\n )\n out[d_str] = nan_factory(d_str)\n continue\n try:\n out[d_str] = compute_fn(model, sample, device=d, **kwargs)\n except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n warnings.warn(f\"{metric_name} benchmark failed on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n except Exception as e:\n warnings.warn(f\"Unexpected error during {metric_name} benchmark on {d}: {e}\")\n out[d_str] = nan_factory(d_str)\n return out"
+ "source": [
+ "#| export\n",
+ "def _bytes_to_mib(n: int) -> float: # number of bytes to convert\n",
+ " \"\"\"Convert bytes to MiB (base-2).\"\"\"\n",
+ " return n / 1024 ** 2\n",
+ "\n",
+ "def _fmt_human(n: int | float) -> str: # number to format\n",
+ " \"\"\"Format large number with K/M/B suffix for human-readable output.\"\"\"\n",
+ " if n >= 1e9: return f\"{n/1e9:.2f}B\"\n",
+ " if n >= 1e6: return f\"{n/1e6:.2f}M\"\n",
+ " if n >= 1e3: return f\"{n/1e3:.2f}K\"\n",
+ " return str(int(n))\n",
+ "\n",
+ "\n",
+ "def _fmt_table(n: int | float) -> str: # number to format\n",
+ " \"\"\"Format number with commas for tabular output.\"\"\"\n",
+ " return f\"{int(n):>12,}\"\n",
+ "\n",
+ "\n",
+ "def _fmt_float(v: float, width: int = 8, decimals: int = 3) -> str: # value, width, decimal places\n",
+ " \"\"\"Format float with specified width and decimal places.\"\"\"\n",
+ " return f\"{v:{width}.{decimals}f}\"\n",
+ "\n",
+ "\n",
+ "def _fmt_macs(v: float) -> str: # MACs value to format\n",
+ " \"\"\"Format MACs with appropriate M/K suffix.\"\"\"\n",
+ " if v >= 1e6: return f\"{v/1e6:8.2f} M\"\n",
+ " if v >= 1e3: return f\"{v/1e3:8.2f} K\"\n",
+ " return f\"{int(v):>10}\"\n",
+ "\n",
+ "\n",
+ "def _section(title: str, width: int = 40) -> str: # section title, total width\n",
+ " \"\"\"Create a section header with box-drawing characters.\"\"\"\n",
+ " return f\"═══ {title} \" + \"═\" * (width - len(title))\n",
+ "\n",
+ "@contextlib.contextmanager\n",
+ "def _device_ctx(dev: str | torch.device): # device string or torch.device\n",
+ " \"\"\"Context manager that validates device availability and yields resolved device.\"\"\"\n",
+ " dev = torch.device(dev)\n",
+ " if dev.type == \"cuda\" and not torch.cuda.is_available():\n",
+ " warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n",
+ " dev = torch.device(\"cpu\")\n",
+ " yield dev\n",
+ "\n",
+ "\n",
+ "@contextlib.contextmanager\n",
+ "def _on_device(model: torch.nn.Module, dev: str | torch.device):\n",
+ " \"\"\"Yield `model` on `dev`, then move it back where it was.\"\"\"\n",
+ " t = next(chain(model.parameters(), model.buffers()), None)\n",
+ " orig = t.device if t is not None else dev # no tensors: nothing to restore\n",
+ " try:\n",
+ " model.to(dev)\n",
+ " yield model\n",
+ " finally:\n",
+ " model.to(orig)\n",
+ "\n",
+ "\n",
+ "def _sync(dev: torch.device) -> None: # device to synchronize\n",
+ " \"\"\"Synchronize CUDA device if applicable.\"\"\"\n",
+ " if dev.type == \"cuda\":\n",
+ " torch.cuda.synchronize(dev)\n",
+ "\n",
+ "\n",
+ "def _default_devices() -> list[str]:\n",
+ " \"\"\"Return default device list: ['cpu'] + ['cuda'] if available.\"\"\"\n",
+ " devices = [\"cpu\"]\n",
+ " if torch.cuda.is_available():\n",
+ " devices.append(\"cuda\")\n",
+ " return devices\n",
+ "\n",
+ "\n",
+ "def _is_quantized(model: torch.nn.Module) -> bool: # model to inspect\n",
+ " \"\"\"True if the model contains PyTorch quantized modules.\n",
+ "\n",
+ " PyTorch's eager/FX quantized ops (fbgemm/qnnpack) are **CPU-only**. Running\n",
+ " them on a CUDA tensor has no registered kernel and *segfaults* the process\n",
+ " rather than raising, so callers must avoid dispatching them off-CPU.\n",
+ " \"\"\"\n",
+ " return any(\"quantized\" in type(m).__module__ for m in model.modules())\n",
+ "\n",
+ "\n",
+ "def _device_supported(\n",
+ " model: torch.nn.Module, # model to run\n",
+ " device: str | torch.device, # target device\n",
+ ") -> bool:\n",
+ " \"\"\"False if `model` cannot execute on `device` (quantized models are CPU-only).\"\"\"\n",
+ " return not (torch.device(device).type != \"cpu\" and _is_quantized(model))\n",
+ "\n",
+ "\n",
+ "def _ensure_device_supported(\n",
+ " model: torch.nn.Module, # model to run\n",
+ " device: str | torch.device, # target device\n",
+ ") -> None:\n",
+ " \"\"\"Raise a *catchable* error instead of letting PyTorch segfault.\n",
+ "\n",
+ " Guards the forward-executing profilers: a quantized model dispatched on a\n",
+ " non-CPU device would crash the interpreter with SIGSEGV, so we fail loudly\n",
+ " (and recoverably) first.\n",
+ " \"\"\"\n",
+ " if not _device_supported(model, device):\n",
+ " raise RuntimeError(\n",
+ " f\"Quantized models run on CPU only; cannot benchmark on '{device}'. \"\n",
+ " \"PyTorch quantized ops (fbgemm/qnnpack) have no CUDA kernels.\"\n",
+ " )\n",
+ "\n",
+ "\n",
+ "def _run_on_devices(\n",
+ " compute_fn, # single-device compute function\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor\n",
+ " devices: list[str | torch.device] | None, # devices to run on (None = default)\n",
+ " nan_factory, # callable(device_str) -> metrics with NaN values\n",
+ " metric_name: str, # for warning messages (e.g., \"Speed\", \"Memory\")\n",
+ " **kwargs,\n",
+ ") -> dict:\n",
+ " \"\"\"Run compute_fn on multiple devices with unified error handling.\n",
+ " \n",
+ " Returns dict mapping device string to metrics (or NaN metrics on failure).\n",
+ " Devices a model cannot run on (e.g. CUDA for a quantized model) are skipped\n",
+ " with NaN metrics — this avoids an uncatchable SIGSEGV from dispatching\n",
+ " CPU-only quantized ops on a CUDA tensor.\n",
+ " \"\"\"\n",
+ " if devices is None:\n",
+ " devices = _default_devices()\n",
+ "\n",
+ " out = {}\n",
+ " for d in devices:\n",
+ " d_str = str(d)\n",
+ " if not _device_supported(model, d):\n",
+ " warnings.warn(\n",
+ " f\"{metric_name} benchmark skipped on {d}: quantized models run on \"\n",
+ " \"CPU only (PyTorch quantized ops have no CUDA kernels).\"\n",
+ " )\n",
+ " out[d_str] = nan_factory(d_str)\n",
+ " continue\n",
+ " try:\n",
+ " out[d_str] = compute_fn(model, sample, device=d, **kwargs)\n",
+ " except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n",
+ " warnings.warn(f\"{metric_name} benchmark failed on {d}: {e}\")\n",
+ " out[d_str] = nan_factory(d_str)\n",
+ " except Exception as e:\n",
+ " warnings.warn(f\"Unexpected error during {metric_name} benchmark on {d}: {e}\")\n",
+ " out[d_str] = nan_factory(d_str)\n",
+ " return out"
+ ]
},
{
"cell_type": "code",
@@ -56,7 +201,61 @@
"id": "1d0944ae",
"metadata": {},
"outputs": [],
- "source": "#| hide\nfrom fastcore.test import *\nimport torch.nn as nn\n\n# _is_quantized: plain fp modules are not quantized\n_fp = nn.Sequential(nn.Conv2d(3, 4, 3), nn.ReLU())\ntest_eq(_is_quantized(_fp), False)\n# fp models can run anywhere; quantized models are CPU-only\ntest_eq(_device_supported(_fp, \"cuda\"), True)\ntest_eq(_device_supported(_fp, \"cpu\"), True)\n_ensure_device_supported(_fp, \"cuda\") # no raise for fp model\n\n# A module whose class lives under a *.quantized.* path is detected as quantized\nimport torch.ao.nn.quantized as nnq\n_q = nn.Sequential(nnq.Linear(4, 2))\ntest_eq(_is_quantized(_q), True)\ntest_eq(_device_supported(_q, \"cpu\"), True)\ntest_eq(_device_supported(_q, \"cuda\"), False)\nwith ExceptionExpected(RuntimeError):\n _ensure_device_supported(_q, \"cuda\")"
+ "source": [
+ "#| hide\n",
+ "from fastcore.test import *\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "# _is_quantized: plain fp modules are not quantized\n",
+ "_fp = nn.Sequential(nn.Conv2d(3, 4, 3), nn.ReLU())\n",
+ "test_eq(_is_quantized(_fp), False)\n",
+ "# fp models can run anywhere; quantized models are CPU-only\n",
+ "test_eq(_device_supported(_fp, \"cuda\"), True)\n",
+ "test_eq(_device_supported(_fp, \"cpu\"), True)\n",
+ "_ensure_device_supported(_fp, \"cuda\") # no raise for fp model\n",
+ "\n",
+ "# A module whose class lives under a *.quantized.* path is detected as quantized\n",
+ "import torch.ao.nn.quantized as nnq\n",
+ "_q = nn.Sequential(nnq.Linear(4, 2))\n",
+ "test_eq(_is_quantized(_q), True)\n",
+ "test_eq(_device_supported(_q, \"cpu\"), True)\n",
+ "test_eq(_device_supported(_q, \"cuda\"), False)\n",
+ "with ExceptionExpected(RuntimeError):\n",
+ " _ensure_device_supported(_q, \"cuda\")\n",
+ "\n",
+ "# _on_device yields the model on the target device and restores the original one\n",
+ "_m = nn.Linear(4, 2)\n",
+ "with _on_device(_m, \"cpu\") as _inner:\n",
+ " assert _inner is _m\n",
+ " test_eq({p.device.type for p in _inner.parameters()}, {\"cpu\"})\n",
+ "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"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "5760b663",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: a CUDA model measured on the CPU comes back on CUDA.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Linear(4, 2).cuda()\n",
+ " with _on_device(_cm, \"cpu\") as _inner:\n",
+ " test_eq({p.device.type for p in _inner.parameters()}, {\"cpu\"})\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+ "\n",
+ " # ...restored even when the body raises\n",
+ " 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\"})"
+ ]
},
{
"cell_type": "markdown",
diff --git a/nbs/index.ipynb b/nbs/index.ipynb
index 9c8be32..a14b667 100644
--- a/nbs/index.ipynb
+++ b/nbs/index.ipynb
@@ -13,32 +13,54 @@
{
"cell_type": "code",
"execution_count": null,
- "id": "389be194",
- "metadata": {},
- "outputs": [],
- "source": [
- "#| default_exp __init__"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "version_cell",
+ "id": "8b6f8c52",
"metadata": {},
"outputs": [],
"source": [
- "#| export\n",
- "__version__ = \"0.1.0\""
+ "from fasterbench.benchmark import benchmark, BenchmarkResult\n",
+ "from fasterbench.size import SizeMetrics, compute_size, get_model_size, get_num_parameters\n",
+ "from fasterbench.speed import (\n",
+ " SpeedMetrics, compute_speed, compute_speed_multi, \n",
+ " sweep_threads, sweep_latency, sweep_batch_sizes,\n",
+ ")\n",
+ "from fasterbench.profiling import LayerProfiler\n",
+ "from fasterbench.compute import ComputeMetrics, compute_compute\n",
+ "from fasterbench.memory import MemoryMetrics, compute_memory, compute_memory_multi\n",
+ "from fasterbench.energy import EnergyMetrics, compute_energy, compute_energy_multi\n",
+ "from fasterbench.roofline import (\n",
+ " HardwarePeaks, RooflinePoint, measure_peaks, clear_peaks_cache, RooflineAnalyzer,\n",
+ ")\n",
+ "from fasterbench.plot import create_radar_plot, SPECS\n",
+ "from fasterbench.utils import parse_metric_value\n",
+ "\n",
+ "__all__ = [\n",
+ " # Main entry point\n",
+ " 'benchmark', 'BenchmarkResult',\n",
+ " # Size\n",
+ " 'SizeMetrics', 'compute_size', 'get_model_size', 'get_num_parameters',\n",
+ " # Speed\n",
+ " 'SpeedMetrics', 'compute_speed', 'compute_speed_multi', \n",
+ " 'sweep_threads', 'sweep_latency', 'sweep_batch_sizes',\n",
+ " # Profiling\n",
+ " 'LayerProfiler',\n",
+ " # Compute\n",
+ " 'ComputeMetrics', 'compute_compute',\n",
+ " # Memory\n",
+ " 'MemoryMetrics', 'compute_memory', 'compute_memory_multi',\n",
+ " # Energy\n",
+ " 'EnergyMetrics', 'compute_energy', 'compute_energy_multi',\n",
+ " # Roofline\n",
+ " 'HardwarePeaks', 'RooflinePoint', 'measure_peaks', 'clear_peaks_cache', 'RooflineAnalyzer',\n",
+ " # Plot\n",
+ " 'create_radar_plot', 'SPECS',\n",
+ " # Report\n",
+ " 'Report', 'ComparisonReport', 'ReportMetricDelta',\n",
+ " # Utils\n",
+ " 'parse_metric_value',\n",
+ "]\n",
+ "from fasterbench.report import Report, ComparisonReport, ReportMetricDelta\n"
]
},
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "8b6f8c52",
- "metadata": {},
- "outputs": [],
- "source": "#| export\nfrom fasterbench.benchmark import benchmark, BenchmarkResult\nfrom fasterbench.size import SizeMetrics, compute_size, get_model_size, get_num_parameters\nfrom fasterbench.speed import (\n SpeedMetrics, compute_speed, compute_speed_multi, \n sweep_threads, sweep_latency, sweep_batch_sizes,\n)\nfrom fasterbench.profiling import LayerProfiler\nfrom fasterbench.compute import ComputeMetrics, compute_compute\nfrom fasterbench.memory import MemoryMetrics, compute_memory, compute_memory_multi\nfrom fasterbench.energy import EnergyMetrics, compute_energy, compute_energy_multi\nfrom fasterbench.roofline import (\n HardwarePeaks, RooflinePoint, measure_peaks, clear_peaks_cache, RooflineAnalyzer,\n)\nfrom fasterbench.plot import create_radar_plot, SPECS\nfrom fasterbench.utils import parse_metric_value\n\n__all__ = [\n # Main entry point\n 'benchmark', 'BenchmarkResult',\n # Size\n 'SizeMetrics', 'compute_size', 'get_model_size', 'get_num_parameters',\n # Speed\n 'SpeedMetrics', 'compute_speed', 'compute_speed_multi', \n 'sweep_threads', 'sweep_latency', 'sweep_batch_sizes',\n # Profiling\n 'LayerProfiler',\n # Compute\n 'ComputeMetrics', 'compute_compute',\n # Memory\n 'MemoryMetrics', 'compute_memory', 'compute_memory_multi',\n # Energy\n 'EnergyMetrics', 'compute_energy', 'compute_energy_multi',\n # Roofline\n 'HardwarePeaks', 'RooflinePoint', 'measure_peaks', 'clear_peaks_cache', 'RooflineAnalyzer',\n # Plot\n 'create_radar_plot', 'SPECS',\n # Report\n 'Report', 'ComparisonReport', 'ReportMetricDelta',\n # Utils\n 'parse_metric_value',\n]\nfrom fasterbench.report import Report, ComparisonReport, ReportMetricDelta\n"
- },
{
"cell_type": "markdown",
"id": "520239d3",
diff --git a/nbs/metrics/compute.ipynb b/nbs/metrics/compute.ipynb
index 3df3272..bd5abc3 100644
--- a/nbs/metrics/compute.ipynb
+++ b/nbs/metrics/compute.ipynb
@@ -34,7 +34,29 @@
"id": "0091d170",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nimport math\nimport warnings\nfrom dataclasses import dataclass\n\nimport torch\nimport torch.nn as nn\n\nfrom fasterbench.core import _is_quantized\n\ntry:\n from thop import profile as _thop_profile\nexcept ImportError:\n _thop_profile = None\n\ntry:\n from torchprofile import profile_macs as _profile_macs\nexcept ImportError:\n _profile_macs = None"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "import math\n",
+ "import warnings\n",
+ "from dataclasses import dataclass\n",
+ "\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "from fasterbench.core import _is_quantized\n",
+ "\n",
+ "try:\n",
+ " from thop import profile as _thop_profile\n",
+ "except ImportError:\n",
+ " _thop_profile = None\n",
+ "\n",
+ "try:\n",
+ " from torchprofile import profile_macs as _profile_macs\n",
+ "except ImportError:\n",
+ " _profile_macs = None"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +64,107 @@
"id": "09401efb",
"metadata": {},
"outputs": [],
- "source": "#| export\n@dataclass(slots=True)\nclass ComputeMetrics:\n \"\"\"MACs (Multiply-Accumulate operations) in millions.\"\"\"\n macs_m: float | None # None if unavailable\n\n @property\n def macs_available(self) -> bool:\n \"\"\"Check if MACs measurement succeeded.\"\"\"\n return self.macs_m is not None\n\n def as_dict(self) -> dict[str, float]:\n return {\n \"macs_m\": self.macs_m if self.macs_m is not None else float(\"nan\"),\n }\n\n\n#| export\ndef _count_macs_hooks(\n model: nn.Module, # model to profile (run on CPU)\n sample: torch.Tensor, # input tensor (with batch dimension)\n) -> int:\n \"\"\"Count Conv2d/Linear MACs via forward hooks (fp *and* quantized modules).\n\n thop has no handlers for quantized ops (it reports ~0), so we count manually.\n Quantization does not change the number of multiply-accumulates: an int8 conv\n does the same MACs as its fp32 twin. Uses module attributes (`out_channels`,\n `kernel_size`, `in_features`, ...) which both fp and quantized modules expose,\n and reads output spatial dims from the produced tensor.\n \"\"\"\n macs = 0\n handles = []\n\n def conv_hook(m, inp, out):\n nonlocal macs\n out_spatial = math.prod(out.shape[2:]) # out_H * out_W\n kernel = math.prod(m.kernel_size) # kH * kW\n macs += out.shape[0] * m.out_channels * (m.in_channels // m.groups) * kernel * out_spatial\n\n def linear_hook(m, inp, out):\n nonlocal macs\n leading = out.numel() // m.out_features # batch/token dims\n macs += leading * m.in_features * m.out_features\n\n for m in model.modules():\n if hasattr(m, \"kernel_size\") and hasattr(m, \"out_channels\") and hasattr(m, \"groups\"):\n handles.append(m.register_forward_hook(conv_hook))\n elif hasattr(m, \"in_features\") and hasattr(m, \"out_features\"):\n handles.append(m.register_forward_hook(linear_hook))\n\n try:\n with torch.no_grad():\n model(sample)\n finally:\n for h in handles:\n h.remove()\n return macs\n\n\n#| export\ndef compute_compute(\n model: nn.Module, # model to analyze\n sample: torch.Tensor, # input tensor (with batch dimension)\n) -> ComputeMetrics:\n \"\"\"Compute MACs for a single forward pass.\"\"\"\n try:\n model_device = next(model.parameters()).device\n except StopIteration:\n model_device = torch.device(\"cpu\")\n\n if sample.device != model_device:\n sample = sample.to(model_device)\n\n macs_m: float | None = None\n\n # Quantized models: thop has no quantized-op handlers (returns ~0). Count the\n # Conv2d/Linear MACs manually — quantization preserves the MAC count.\n if _is_quantized(model):\n try:\n macs_m = round(_count_macs_hooks(model, sample) / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"quantized MAC counting failed: {e}\")\n return ComputeMetrics(macs_m=macs_m)\n\n if _thop_profile is not None:\n try:\n mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False)\n macs_m = round(mac_raw / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"thop failed: {e}\")\n elif _profile_macs is not None:\n try:\n macs_m = round(_profile_macs(model, sample) / 1e6, 3)\n except Exception as e:\n warnings.warn(f\"torchprofile failed: {e}\")\n else:\n warnings.warn(\"No MAC-counting backend available – skipping MACs\")\n\n return ComputeMetrics(macs_m=macs_m)"
+ "source": [
+ "#| export\n",
+ "@dataclass(slots=True)\n",
+ "class ComputeMetrics:\n",
+ " \"\"\"MACs (Multiply-Accumulate operations) in millions.\"\"\"\n",
+ " macs_m: float | None # None if unavailable\n",
+ "\n",
+ " @property\n",
+ " def macs_available(self) -> bool:\n",
+ " \"\"\"Check if MACs measurement succeeded.\"\"\"\n",
+ " return self.macs_m is not None\n",
+ "\n",
+ " def as_dict(self) -> dict[str, float]:\n",
+ " return {\n",
+ " \"macs_m\": self.macs_m if self.macs_m is not None else float(\"nan\"),\n",
+ " }\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _count_macs_hooks(\n",
+ " model: nn.Module, # model to profile (run on CPU)\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ ") -> int:\n",
+ " \"\"\"Count Conv2d/Linear MACs via forward hooks (fp *and* quantized modules).\n",
+ "\n",
+ " thop has no handlers for quantized ops (it reports ~0), so we count manually.\n",
+ " Quantization does not change the number of multiply-accumulates: an int8 conv\n",
+ " does the same MACs as its fp32 twin. Uses module attributes (`out_channels`,\n",
+ " `kernel_size`, `in_features`, ...) which both fp and quantized modules expose,\n",
+ " and reads output spatial dims from the produced tensor.\n",
+ " \"\"\"\n",
+ " macs = 0\n",
+ " handles = []\n",
+ "\n",
+ " def conv_hook(m, inp, out):\n",
+ " nonlocal macs\n",
+ " out_spatial = math.prod(out.shape[2:]) # out_H * out_W\n",
+ " kernel = math.prod(m.kernel_size) # kH * kW\n",
+ " macs += out.shape[0] * m.out_channels * (m.in_channels // m.groups) * kernel * out_spatial\n",
+ "\n",
+ " def linear_hook(m, inp, out):\n",
+ " nonlocal macs\n",
+ " leading = out.numel() // m.out_features # batch/token dims\n",
+ " macs += leading * m.in_features * m.out_features\n",
+ "\n",
+ " for m in model.modules():\n",
+ " if hasattr(m, \"kernel_size\") and hasattr(m, \"out_channels\") and hasattr(m, \"groups\"):\n",
+ " handles.append(m.register_forward_hook(conv_hook))\n",
+ " elif hasattr(m, \"in_features\") and hasattr(m, \"out_features\"):\n",
+ " handles.append(m.register_forward_hook(linear_hook))\n",
+ "\n",
+ " try:\n",
+ " with torch.no_grad():\n",
+ " model(sample)\n",
+ " finally:\n",
+ " for h in handles:\n",
+ " h.remove()\n",
+ " return macs\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_compute(\n",
+ " model: nn.Module, # model to analyze\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ ") -> ComputeMetrics:\n",
+ " \"\"\"Compute MACs for a single forward pass.\"\"\"\n",
+ " try:\n",
+ " model_device = next(model.parameters()).device\n",
+ " except StopIteration:\n",
+ " model_device = torch.device(\"cpu\")\n",
+ "\n",
+ " if sample.device != model_device:\n",
+ " sample = sample.to(model_device)\n",
+ "\n",
+ " macs_m: float | None = None\n",
+ "\n",
+ " # Quantized models: thop has no quantized-op handlers (returns ~0). Count the\n",
+ " # Conv2d/Linear MACs manually — quantization preserves the MAC count.\n",
+ " if _is_quantized(model):\n",
+ " try:\n",
+ " macs_m = round(_count_macs_hooks(model, sample) / 1e6, 3)\n",
+ " except Exception as e:\n",
+ " warnings.warn(f\"quantized MAC counting failed: {e}\")\n",
+ " return ComputeMetrics(macs_m=macs_m)\n",
+ "\n",
+ " if _thop_profile is not None:\n",
+ " try:\n",
+ " mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False)\n",
+ " macs_m = round(mac_raw / 1e6, 3)\n",
+ " except Exception as e:\n",
+ " warnings.warn(f\"thop failed: {e}\")\n",
+ " elif _profile_macs is not None:\n",
+ " try:\n",
+ " macs_m = round(_profile_macs(model, sample) / 1e6, 3)\n",
+ " except Exception as e:\n",
+ " warnings.warn(f\"torchprofile failed: {e}\")\n",
+ " else:\n",
+ " warnings.warn(\"No MAC-counting backend available – skipping MACs\")\n",
+ "\n",
+ " return ComputeMetrics(macs_m=macs_m)"
+ ]
},
{
"cell_type": "code",
@@ -88,7 +210,46 @@
"id": "b7d64a51",
"metadata": {},
"outputs": [],
- "source": "#| hide\n# Quantization preserves the *number* of MACs (int8 conv/linear do the same\n# multiply-accumulates as their fp32 twin) — only the ops get cheaper. But thop\n# has no quantized-op handlers, so it reports ~0. Verify the manual hook counter\n# recovers a MAC count ~equal to the fp32 twin's thop count.\nfrom torch.ao.quantization import get_default_qconfig, QConfigMapping\nfrom torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\nfrom fasterbench.core import _is_quantized\n\nclass _MacNet(nn.Module):\n \"Conv-dominated net (linear/pool MACs negligible), statically quantizable via FX.\"\n def __init__(self):\n super().__init__()\n self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n def forward(self, x):\n x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n return self.fc(self.flat(self.pool(x)))\n\n_fp = _MacNet().eval()\n_x = torch.randn(1, 3, 16, 16)\n_fp_macs = compute_compute(_fp, _x).macs_m # fp32 baseline (thop path)\nassert _fp_macs is not None and _fp_macs > 0\n\ntorch.backends.quantized.engine = \"x86\"\n_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n_prep = prepare_fx(_MacNet().eval(), _qmap, example_inputs=(_x,))\n_prep(_x) # calibrate\n_qm = convert_fx(_prep)\nassert _is_quantized(_qm)\n\n# BEFORE the fix, thop reports ~0 MACs for a fully quantized model.\n_thop_raw, _ = _thop_profile(_qm, inputs=(_x,), verbose=False)\nassert _thop_raw / 1e6 < 0.1 * _fp_macs # thop is effectively blind here\n\n_q = compute_compute(_qm, _x) # AFTER the fix (manual counter)\nassert _q.macs_m is not None and _q.macs_m > 0 # zero is gone\nassert abs(_q.macs_m - _fp_macs) <= 0.05 * _fp_macs # within ~5% of fp32 MACs"
+ "source": [
+ "#| hide\n",
+ "# Quantization preserves the *number* of MACs (int8 conv/linear do the same\n",
+ "# multiply-accumulates as their fp32 twin) — only the ops get cheaper. But thop\n",
+ "# has no quantized-op handlers, so it reports ~0. Verify the manual hook counter\n",
+ "# recovers a MAC count ~equal to the fp32 twin's thop count.\n",
+ "from torch.ao.quantization import get_default_qconfig, QConfigMapping\n",
+ "from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\n",
+ "from fasterbench.core import _is_quantized\n",
+ "\n",
+ "class _MacNet(nn.Module):\n",
+ " \"Conv-dominated net (linear/pool MACs negligible), statically quantizable via FX.\"\n",
+ " def __init__(self):\n",
+ " super().__init__()\n",
+ " self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n",
+ " self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n",
+ " def forward(self, x):\n",
+ " x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n",
+ " return self.fc(self.flat(self.pool(x)))\n",
+ "\n",
+ "_fp = _MacNet().eval()\n",
+ "_x = torch.randn(1, 3, 16, 16)\n",
+ "_fp_macs = compute_compute(_fp, _x).macs_m # fp32 baseline (thop path)\n",
+ "assert _fp_macs is not None and _fp_macs > 0\n",
+ "\n",
+ "torch.backends.quantized.engine = \"x86\"\n",
+ "_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n",
+ "_prep = prepare_fx(_MacNet().eval(), _qmap, example_inputs=(_x,))\n",
+ "_prep(_x) # calibrate\n",
+ "_qm = convert_fx(_prep)\n",
+ "assert _is_quantized(_qm)\n",
+ "\n",
+ "# BEFORE the fix, thop reports ~0 MACs for a fully quantized model.\n",
+ "_thop_raw, _ = _thop_profile(_qm, inputs=(_x,), verbose=False)\n",
+ "assert _thop_raw / 1e6 < 0.1 * _fp_macs # thop is effectively blind here\n",
+ "\n",
+ "_q = compute_compute(_qm, _x) # AFTER the fix (manual counter)\n",
+ "assert _q.macs_m is not None and _q.macs_m > 0 # zero is gone\n",
+ "assert abs(_q.macs_m - _fp_macs) <= 0.05 * _fp_macs # within ~5% of fp32 MACs"
+ ]
},
{
"cell_type": "markdown",
diff --git a/nbs/metrics/energy.ipynb b/nbs/metrics/energy.ipynb
index 1366ee8..5c984f9 100644
--- a/nbs/metrics/energy.ipynb
+++ b/nbs/metrics/energy.ipynb
@@ -34,7 +34,25 @@
"id": "d27f26a4",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\nimport logging, os, time, warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\n\nimport torch\n\nfrom fasterbench.core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported\n\ntry:\n from codecarbon import EmissionsTracker, OfflineEmissionsTracker\n # Suppress codecarbon's verbose logging BEFORE any tracker is constructed,\n # because codecarbon logs messages in __init__ before applying log_level.\n logging.getLogger(\"codecarbon\").setLevel(logging.CRITICAL)\nexcept ImportError:\n EmissionsTracker = OfflineEmissionsTracker = None"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "import logging, os, time, warnings\n",
+ "from dataclasses import dataclass, asdict\n",
+ "from typing import Sequence\n",
+ "\n",
+ "import torch\n",
+ "\n",
+ "from fasterbench.core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported\n",
+ "\n",
+ "try:\n",
+ " from codecarbon import EmissionsTracker, OfflineEmissionsTracker\n",
+ " # Suppress codecarbon's verbose logging BEFORE any tracker is constructed,\n",
+ " # because codecarbon logs messages in __init__ before applying log_level.\n",
+ " logging.getLogger(\"codecarbon\").setLevel(logging.CRITICAL)\n",
+ "except ImportError:\n",
+ " EmissionsTracker = OfflineEmissionsTracker = None"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +60,131 @@
"id": "58dcf143",
"metadata": {},
"outputs": [],
- "source": "#| export\n@dataclass(slots=True)\nclass EnergyMetrics:\n \"\"\"Energy consumption and carbon footprint metrics.\"\"\"\n mean_watts: float # average power during measurement\n energy_wh: float # Wh per inference\n co2_eq_g: float # g CO₂-eq per inference\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef _nan_energy_metrics(device: str) -> EnergyMetrics: # device string (unused, for consistent signature)\n \"\"\"Create EnergyMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return EnergyMetrics(nan, nan, nan)\n\n\n#| export\ndef _clear_stale_codecarbon_lock() -> None:\n \"\"\"Remove stale codecarbon lock file if the owning process no longer exists.\"\"\"\n import tempfile\n lock_path = os.path.join(tempfile.gettempdir(), \".codecarbon.lock\")\n if not os.path.exists(lock_path):\n return\n try:\n # Read the PID from the lock file (codecarbon writes its PID there)\n with open(lock_path) as f:\n content = f.read().strip()\n if content:\n pid = int(content)\n os.kill(pid, 0) # Check if process exists (signal 0 = no-op)\n # Process exists — lock is valid, don't remove\n return\n except (ValueError, ProcessLookupError, PermissionError, OSError):\n pass # PID invalid or process dead — lock is stale\n try:\n os.remove(lock_path)\n except OSError:\n pass\n\n\n#| export\n@torch.no_grad()\ndef compute_energy(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n offline: bool = True, # use offline emissions tracker\n country_iso: str | None = None, # country ISO code for carbon intensity\n measure_secs: int = 1, # power sampling interval\n) -> EnergyMetrics:\n \"\"\"Measure power consumption and carbon footprint using codecarbon.\"\"\"\n if EmissionsTracker is None:\n warnings.warn(\"codecarbon not installed – returning NaNs\")\n return _nan_energy_metrics(str(device))\n\n _clear_stale_codecarbon_lock()\n\n Tracker = OfflineEmissionsTracker if offline else EmissionsTracker\n tracker = Tracker(\n project_name=\"fasterbench\",\n country_iso_code=(country_iso or os.getenv(\"NNBENCH_ISO\", \"USA\")),\n measure_power_secs=measure_secs,\n save_to_file=False,\n log_level=\"critical\",\n )\n\n with _device_ctx(device) as dev:\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model = model.eval().to(dev)\n sample = sample.to(dev, non_blocking=True)\n\n for _ in range(warmup):\n model(sample)\n _sync(dev)\n\n tracker.start()\n try:\n t0 = time.perf_counter()\n for _ in range(steps):\n model(sample)\n _sync(dev)\n finally:\n tracker.stop()\n dur_s = time.perf_counter() - t0\n\n # codecarbon silently fails if another instance is running,\n # leaving final_emissions_data as None\n if tracker.final_emissions_data is None:\n warnings.warn(\"codecarbon tracker did not collect data (another instance may be running)\")\n return _nan_energy_metrics(str(device))\n\n ene_kwh = tracker.final_emissions_data.energy_consumed\n co2_kg = tracker.final_emissions\n mean_w = (ene_kwh * 3600_000) / dur_s\n\n return EnergyMetrics(\n mean_watts=mean_w,\n energy_wh=(ene_kwh * 1_000) / steps,\n co2_eq_g=(co2_kg * 1_000) / steps,\n )\n\n\n#| export\ndef compute_energy_multi(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n **kwargs,\n) -> dict[str, EnergyMetrics]:\n \"\"\"Measure energy on multiple devices.\"\"\"\n return _run_on_devices(\n compute_energy, model, sample, devices,\n nan_factory=_nan_energy_metrics,\n metric_name=\"Energy\",\n **kwargs\n )"
+ "source": [
+ "#| export\n",
+ "@dataclass(slots=True)\n",
+ "class EnergyMetrics:\n",
+ " \"\"\"Energy consumption and carbon footprint metrics.\"\"\"\n",
+ " mean_watts: float # average power during measurement\n",
+ " energy_wh: float # Wh per inference\n",
+ " co2_eq_g: float # g CO₂-eq per inference\n",
+ "\n",
+ " def as_dict(self) -> dict[str, float]:\n",
+ " return asdict(self)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _nan_energy_metrics(device: str) -> EnergyMetrics: # device string (unused, for consistent signature)\n",
+ " \"\"\"Create EnergyMetrics with NaN values for failed benchmarks.\"\"\"\n",
+ " nan = float(\"nan\")\n",
+ " return EnergyMetrics(nan, nan, nan)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _clear_stale_codecarbon_lock() -> None:\n",
+ " \"\"\"Remove stale codecarbon lock file if the owning process no longer exists.\"\"\"\n",
+ " import tempfile\n",
+ " lock_path = os.path.join(tempfile.gettempdir(), \".codecarbon.lock\")\n",
+ " if not os.path.exists(lock_path):\n",
+ " return\n",
+ " try:\n",
+ " # Read the PID from the lock file (codecarbon writes its PID there)\n",
+ " with open(lock_path) as f:\n",
+ " content = f.read().strip()\n",
+ " if content:\n",
+ " pid = int(content)\n",
+ " os.kill(pid, 0) # Check if process exists (signal 0 = no-op)\n",
+ " # Process exists — lock is valid, don't remove\n",
+ " return\n",
+ " except (ValueError, ProcessLookupError, PermissionError, OSError):\n",
+ " pass # PID invalid or process dead — lock is stale\n",
+ " try:\n",
+ " os.remove(lock_path)\n",
+ " except OSError:\n",
+ " pass\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "@torch.no_grad()\n",
+ "def compute_energy(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " device: str | torch.device = \"cpu\", # device to run on\n",
+ " warmup: int = 20, # warmup iterations\n",
+ " steps: int = 100, # measurement iterations\n",
+ " offline: bool = True, # use offline emissions tracker\n",
+ " country_iso: str | None = None, # country ISO code for carbon intensity\n",
+ " measure_secs: int = 1, # power sampling interval\n",
+ ") -> EnergyMetrics:\n",
+ " \"\"\"Measure power consumption and carbon footprint using codecarbon.\"\"\"\n",
+ " if EmissionsTracker is None:\n",
+ " warnings.warn(\"codecarbon not installed – returning NaNs\")\n",
+ " return _nan_energy_metrics(str(device))\n",
+ "\n",
+ " _clear_stale_codecarbon_lock()\n",
+ "\n",
+ " Tracker = OfflineEmissionsTracker if offline else EmissionsTracker\n",
+ " tracker = Tracker(\n",
+ " project_name=\"fasterbench\",\n",
+ " country_iso_code=(country_iso or os.getenv(\"NNBENCH_ISO\", \"USA\")),\n",
+ " measure_power_secs=measure_secs,\n",
+ " save_to_file=False,\n",
+ " log_level=\"critical\",\n",
+ " )\n",
+ "\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",
+ " sample = sample.to(dev, non_blocking=True)\n",
+ "\n",
+ " for _ in range(warmup):\n",
+ " model(sample)\n",
+ " _sync(dev)\n",
+ "\n",
+ " tracker.start()\n",
+ " try:\n",
+ " t0 = time.perf_counter()\n",
+ " for _ in range(steps):\n",
+ " model(sample)\n",
+ " _sync(dev)\n",
+ " finally:\n",
+ " tracker.stop()\n",
+ " dur_s = time.perf_counter() - t0\n",
+ "\n",
+ " # codecarbon silently fails if another instance is running,\n",
+ " # leaving final_emissions_data as None\n",
+ " if tracker.final_emissions_data is None:\n",
+ " warnings.warn(\"codecarbon tracker did not collect data (another instance may be running)\")\n",
+ " return _nan_energy_metrics(str(device))\n",
+ "\n",
+ " ene_kwh = tracker.final_emissions_data.energy_consumed\n",
+ " co2_kg = tracker.final_emissions\n",
+ " mean_w = (ene_kwh * 3600_000) / dur_s\n",
+ "\n",
+ " return EnergyMetrics(\n",
+ " mean_watts=mean_w,\n",
+ " energy_wh=(ene_kwh * 1_000) / steps,\n",
+ " co2_eq_g=(co2_kg * 1_000) / steps,\n",
+ " )\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_energy_multi(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n",
+ " **kwargs,\n",
+ ") -> dict[str, EnergyMetrics]:\n",
+ " \"\"\"Measure energy on multiple devices.\"\"\"\n",
+ " return _run_on_devices(\n",
+ " compute_energy, model, sample, devices,\n",
+ " nan_factory=_nan_energy_metrics,\n",
+ " metric_name=\"Energy\",\n",
+ " **kwargs\n",
+ " )"
+ ]
},
{
"cell_type": "code",
@@ -82,6 +224,30 @@
"show_doc(compute_energy_multi)"
]
},
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "9ac48c40",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Runs codecarbon (and needs a second device): measuring leaves the model where it was.\n",
+ "from fastcore.test import *\n",
+ "import torch.nn as nn\n",
+ "\n",
+ "_m = nn.Linear(10, 5)\n",
+ "_x = torch.randn(1, 10)\n",
+ "compute_energy(_m, _x, warmup=1, steps=2)\n",
+ "test_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})\n",
+ "\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Linear(10, 5).cuda()\n",
+ " compute_energy(_cm, _x, device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
+ ]
+ },
{
"cell_type": "markdown",
"id": "see_also",
diff --git a/nbs/metrics/memory.ipynb b/nbs/metrics/memory.ipynb
index 01f6be5..2b5cce5 100644
--- a/nbs/metrics/memory.ipynb
+++ b/nbs/metrics/memory.ipynb
@@ -34,7 +34,23 @@
"id": "16cd91b6",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nimport warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Sequence\nfrom fasterbench.core import _bytes_to_mib, _run_on_devices, _ensure_device_supported\n\nimport numpy as np\nimport torch\n\ntry:\n import psutil\nexcept ImportError:\n psutil = None"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "import warnings\n",
+ "from dataclasses import dataclass, asdict\n",
+ "from typing import Sequence\n",
+ "from fasterbench.core import _bytes_to_mib, _on_device, _run_on_devices, _ensure_device_supported\n",
+ "\n",
+ "import numpy as np\n",
+ "import torch\n",
+ "\n",
+ "try:\n",
+ " import psutil\n",
+ "except ImportError:\n",
+ " psutil = None"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +58,138 @@
"id": "08a0264a",
"metadata": {},
"outputs": [],
- "source": "#| export\n@dataclass(slots=True)\nclass MemoryMetrics:\n \"\"\"Memory usage metrics for a single device.\"\"\"\n avg_mib: float\n peak_mib: float\n reserved_mib: float # GPU-only (NaN for CPU)\n device: str\n\n def as_dict(self) -> dict[str, float | str]:\n return asdict(self)\n\n\n#| export\ndef _nan_memory_metrics(device: str) -> MemoryMetrics: # device string\n \"\"\"Create MemoryMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return MemoryMetrics(nan, nan, nan, device)\n\n\n#| export\ndef _gpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure GPU memory usage.\"\"\"\n dev = torch.device(\"cuda\")\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model = model.eval().to(dev)\n sample = sample.to(dev)\n\n for _ in range(warmup):\n model(sample)\n torch.cuda.synchronize(dev)\n\n alloc, alloc_peak, reserv_peak = [], [], []\n for _ in range(steps):\n torch.cuda.reset_peak_memory_stats(dev)\n model(sample)\n torch.cuda.synchronize(dev)\n alloc.append(torch.cuda.memory_allocated(dev))\n alloc_peak.append(torch.cuda.max_memory_allocated(dev))\n reserv_peak.append(torch.cuda.max_memory_reserved(dev))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(alloc))),\n peak_mib=_bytes_to_mib(float(np.mean(alloc_peak))),\n reserved_mib=_bytes_to_mib(float(np.mean(reserv_peak))),\n device=\"cuda\",\n )\n\n\n#| export\ndef _cpu_metrics(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n warmup: int, # warmup iterations\n steps: int, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure CPU memory usage via psutil.\"\"\"\n if psutil is None:\n warnings.warn(\"psutil not available – returning NaNs for CPU memory\")\n return _nan_memory_metrics(\"cpu\")\n\n proc = psutil.Process()\n model = model.eval().cpu()\n sample = sample.cpu()\n\n rss0 = proc.memory_info().rss\n\n for _ in range(warmup):\n model(sample)\n\n diffs: list[int] = []\n peaks: list[int] = []\n for _ in range(steps):\n rss_before = proc.memory_info().rss\n model(sample)\n rss_after = proc.memory_info().rss\n diffs.append(max(0, rss_after - rss_before))\n peaks.append(max(0, rss_after - rss0))\n\n return MemoryMetrics(\n avg_mib=_bytes_to_mib(float(np.mean(diffs))),\n peak_mib=_bytes_to_mib(float(np.max(peaks))),\n reserved_mib=float(\"nan\"),\n device=\"cpu\",\n )\n\n\n#| export\ndef compute_memory(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> MemoryMetrics:\n \"\"\"Measure memory usage on specified device.\"\"\"\n device_str = str(device)\n if device_str == \"cuda\" or (device_str == \"auto\" and torch.cuda.is_available()):\n if not torch.cuda.is_available():\n warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _gpu_metrics(model, sample, warmup=warmup, steps=steps)\n return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n\n\n#| export\ndef compute_memory_multi(\n model: torch.nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n warmup: int = 10, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> dict[str, MemoryMetrics]:\n \"\"\"Measure memory on multiple devices.\"\"\"\n return _run_on_devices(\n compute_memory, model, sample, devices,\n nan_factory=_nan_memory_metrics,\n metric_name=\"Memory\",\n warmup=warmup,\n steps=steps,\n )"
+ "source": [
+ "#| export\n",
+ "@dataclass(slots=True)\n",
+ "class MemoryMetrics:\n",
+ " \"\"\"Memory usage metrics for a single device.\"\"\"\n",
+ " avg_mib: float\n",
+ " peak_mib: float\n",
+ " reserved_mib: float # GPU-only (NaN for CPU)\n",
+ " device: str\n",
+ "\n",
+ " def as_dict(self) -> dict[str, float | str]:\n",
+ " return asdict(self)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _nan_memory_metrics(device: str) -> MemoryMetrics: # device string\n",
+ " \"\"\"Create MemoryMetrics with NaN values for failed benchmarks.\"\"\"\n",
+ " nan = float(\"nan\")\n",
+ " return MemoryMetrics(nan, nan, nan, device)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _gpu_metrics(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " warmup: int, # warmup iterations\n",
+ " steps: int, # measurement iterations\n",
+ ") -> MemoryMetrics:\n",
+ " \"\"\"Measure GPU memory usage.\"\"\"\n",
+ " dev = torch.device(\"cuda\")\n",
+ " _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n",
+ " with _on_device(model.eval(), dev) as model:\n",
+ " sample = sample.to(dev)\n",
+ "\n",
+ " for _ in range(warmup):\n",
+ " model(sample)\n",
+ " torch.cuda.synchronize(dev)\n",
+ "\n",
+ " alloc, alloc_peak, reserv_peak = [], [], []\n",
+ " for _ in range(steps):\n",
+ " torch.cuda.reset_peak_memory_stats(dev)\n",
+ " model(sample)\n",
+ " torch.cuda.synchronize(dev)\n",
+ " alloc.append(torch.cuda.memory_allocated(dev))\n",
+ " alloc_peak.append(torch.cuda.max_memory_allocated(dev))\n",
+ " reserv_peak.append(torch.cuda.max_memory_reserved(dev))\n",
+ "\n",
+ " return MemoryMetrics(\n",
+ " avg_mib=_bytes_to_mib(float(np.mean(alloc))),\n",
+ " peak_mib=_bytes_to_mib(float(np.mean(alloc_peak))),\n",
+ " reserved_mib=_bytes_to_mib(float(np.mean(reserv_peak))),\n",
+ " device=\"cuda\",\n",
+ " )\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _cpu_metrics(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " warmup: int, # warmup iterations\n",
+ " steps: int, # measurement iterations\n",
+ ") -> MemoryMetrics:\n",
+ " \"\"\"Measure CPU memory usage via psutil.\"\"\"\n",
+ " if psutil is None:\n",
+ " warnings.warn(\"psutil not available – returning NaNs for CPU memory\")\n",
+ " return _nan_memory_metrics(\"cpu\")\n",
+ "\n",
+ " proc = psutil.Process()\n",
+ " with _on_device(model.eval(), \"cpu\") as model:\n",
+ " sample = sample.cpu()\n",
+ "\n",
+ " rss0 = proc.memory_info().rss\n",
+ "\n",
+ " for _ in range(warmup):\n",
+ " model(sample)\n",
+ "\n",
+ " diffs: list[int] = []\n",
+ " peaks: list[int] = []\n",
+ " for _ in range(steps):\n",
+ " rss_before = proc.memory_info().rss\n",
+ " model(sample)\n",
+ " rss_after = proc.memory_info().rss\n",
+ " diffs.append(max(0, rss_after - rss_before))\n",
+ " peaks.append(max(0, rss_after - rss0))\n",
+ "\n",
+ " return MemoryMetrics(\n",
+ " avg_mib=_bytes_to_mib(float(np.mean(diffs))),\n",
+ " peak_mib=_bytes_to_mib(float(np.max(peaks))),\n",
+ " reserved_mib=float(\"nan\"),\n",
+ " device=\"cpu\",\n",
+ " )\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_memory(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " device: str | torch.device = \"cpu\", # device to run on\n",
+ " warmup: int = 10, # warmup iterations\n",
+ " steps: int = 100, # measurement iterations\n",
+ ") -> MemoryMetrics:\n",
+ " \"\"\"Measure memory usage on specified device.\"\"\"\n",
+ " device_str = str(device)\n",
+ " if device_str == \"cuda\" or (device_str == \"auto\" and torch.cuda.is_available()):\n",
+ " if not torch.cuda.is_available():\n",
+ " warnings.warn(\"CUDA requested but not available – falling back to CPU\")\n",
+ " return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n",
+ " return _gpu_metrics(model, sample, warmup=warmup, steps=steps)\n",
+ " return _cpu_metrics(model, sample, warmup=warmup, steps=steps)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_memory_multi(\n",
+ " model: torch.nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda)\n",
+ " warmup: int = 10, # warmup iterations\n",
+ " steps: int = 100, # measurement iterations\n",
+ ") -> dict[str, MemoryMetrics]:\n",
+ " \"\"\"Measure memory on multiple devices.\"\"\"\n",
+ " return _run_on_devices(\n",
+ " compute_memory, model, sample, devices,\n",
+ " nan_factory=_nan_memory_metrics,\n",
+ " metric_name=\"Memory\",\n",
+ " warmup=warmup,\n",
+ " steps=steps,\n",
+ " )"
+ ]
},
{
"cell_type": "code",
@@ -89,7 +236,27 @@
"metadata": {},
"outputs": [],
"source": [
- "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_mem = compute_memory(_m, _x, warmup=2, steps=5)\nassert isinstance(_mem, MemoryMetrics)\nassert _mem.avg_mib >= 0"
+ "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_mem = compute_memory(_m, _x, warmup=2, steps=5)\nassert isinstance(_mem, MemoryMetrics)\nassert _mem.avg_mib >= 0\n\n# measuring leaves the model on the device it came in on\ntest_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "ec932def",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: a CPU memory measurement must not park a CUDA model on the CPU.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Linear(10, 5).cuda()\n",
+ " compute_memory(_cm, _x, device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+ " # ...and the CUDA leg leaves a CPU model on the CPU\n",
+ " _pm = nn.Linear(10, 5)\n",
+ " compute_memory(_pm, _x, device=\"cuda\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _pm.parameters()}, {\"cpu\"})"
]
},
{
diff --git a/nbs/metrics/size.ipynb b/nbs/metrics/size.ipynb
index b726a87..cd09fb7 100644
--- a/nbs/metrics/size.ipynb
+++ b/nbs/metrics/size.ipynb
@@ -34,7 +34,15 @@
"id": "4c37777e",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nfrom fasterbench.core import _bytes_to_mib, _is_quantized\nimport torch\nimport io\nfrom dataclasses import dataclass, asdict"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "from fasterbench.core import _bytes_to_mib, _is_quantized\n",
+ "import torch\n",
+ "import io\n",
+ "from dataclasses import dataclass, asdict"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +50,84 @@
"id": "daa12120",
"metadata": {},
"outputs": [],
- "source": "#| export\ndef get_model_size(model: torch.nn.Module) -> int: # model to measure\n \"\"\"Return the on-disk size (bytes) of the serialized model.\"\"\"\n buf = io.BytesIO()\n try:\n model.save(buf)\n except Exception:\n torch.save(model.state_dict(), buf)\n return buf.getbuffer().nbytes\n\n\n#| export\ndef _quantized_param_count(model: torch.nn.Module) -> int: # model with quantized modules\n \"\"\"Sum the numel of packed quantized weights/biases (not in `.parameters()`).\n\n Quantized `Conv2d`/`Linear` keep their weights in `_packed_params`, exposed via\n a *callable* `.weight()`/`.bias()` accessor (unlike a plain `nn.Parameter`, which\n is a tensor). Quantization preserves the number of parameters, so counting these\n makes a quantized model's param count match its fp32 twin.\n \"\"\"\n total = 0\n for m in model.modules():\n for name in (\"weight\", \"bias\"):\n accessor = getattr(m, name, None)\n if not callable(accessor): # plain nn.Parameter tensors are not callable\n continue\n try:\n t = accessor()\n if t is not None:\n total += t.numel()\n except Exception:\n pass\n return total\n\n\n#| export\ndef get_num_parameters(\n model: torch.nn.Module, # model to count parameters\n trainable_only: bool = True, # if True, only count trainable parameters\n) -> int:\n \"\"\"Count the number of (optionally trainable) parameters.\n\n Quantized weights live in packed params outside `.parameters()`, so they are\n added separately (regardless of `trainable_only`, as they are never trainable).\n \"\"\"\n if trainable_only:\n n = sum(p.numel() for p in model.parameters() if p.requires_grad)\n else:\n n = sum(p.numel() for p in model.parameters())\n if _is_quantized(model):\n n += _quantized_param_count(model)\n return n\n\n\n#| export\n@dataclass(slots=True)\nclass SizeMetrics:\n \"\"\"Model size metrics: disk size and parameter count.\"\"\"\n disk_bytes: int\n size_mib: float\n num_params: int\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef compute_size(\n model: torch.nn.Module, # model to measure\n *,\n params_count: int | None = None, # pre-computed parameter count (avoids recount)\n) -> SizeMetrics:\n \"\"\"Compute size metrics for a model.\"\"\"\n disk = get_model_size(model)\n params = params_count if params_count is not None else get_num_parameters(model)\n return SizeMetrics(disk_bytes=disk, size_mib=_bytes_to_mib(disk), num_params=params)"
+ "source": [
+ "#| export\n",
+ "def get_model_size(model: torch.nn.Module) -> int: # model to measure\n",
+ " \"\"\"Return the on-disk size (bytes) of the serialized model.\"\"\"\n",
+ " buf = io.BytesIO()\n",
+ " try:\n",
+ " model.save(buf)\n",
+ " except Exception:\n",
+ " torch.save(model.state_dict(), buf)\n",
+ " return buf.getbuffer().nbytes\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _quantized_param_count(model: torch.nn.Module) -> int: # model with quantized modules\n",
+ " \"\"\"Sum the numel of packed quantized weights/biases (not in `.parameters()`).\n",
+ "\n",
+ " Quantized `Conv2d`/`Linear` keep their weights in `_packed_params`, exposed via\n",
+ " a *callable* `.weight()`/`.bias()` accessor (unlike a plain `nn.Parameter`, which\n",
+ " is a tensor). Quantization preserves the number of parameters, so counting these\n",
+ " makes a quantized model's param count match its fp32 twin.\n",
+ " \"\"\"\n",
+ " total = 0\n",
+ " for m in model.modules():\n",
+ " for name in (\"weight\", \"bias\"):\n",
+ " accessor = getattr(m, name, None)\n",
+ " if not callable(accessor): # plain nn.Parameter tensors are not callable\n",
+ " continue\n",
+ " try:\n",
+ " t = accessor()\n",
+ " if t is not None:\n",
+ " total += t.numel()\n",
+ " except Exception:\n",
+ " pass\n",
+ " return total\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def get_num_parameters(\n",
+ " model: torch.nn.Module, # model to count parameters\n",
+ " trainable_only: bool = True, # if True, only count trainable parameters\n",
+ ") -> int:\n",
+ " \"\"\"Count the number of (optionally trainable) parameters.\n",
+ "\n",
+ " Quantized weights live in packed params outside `.parameters()`, so they are\n",
+ " added separately (regardless of `trainable_only`, as they are never trainable).\n",
+ " \"\"\"\n",
+ " if trainable_only:\n",
+ " n = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+ " else:\n",
+ " n = sum(p.numel() for p in model.parameters())\n",
+ " if _is_quantized(model):\n",
+ " n += _quantized_param_count(model)\n",
+ " return n\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "@dataclass(slots=True)\n",
+ "class SizeMetrics:\n",
+ " \"\"\"Model size metrics: disk size and parameter count.\"\"\"\n",
+ " disk_bytes: int\n",
+ " size_mib: float\n",
+ " num_params: int\n",
+ "\n",
+ " def as_dict(self) -> dict[str, float]:\n",
+ " return asdict(self)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_size(\n",
+ " model: torch.nn.Module, # model to measure\n",
+ " *,\n",
+ " params_count: int | None = None, # pre-computed parameter count (avoids recount)\n",
+ ") -> SizeMetrics:\n",
+ " \"\"\"Compute size metrics for a model.\"\"\"\n",
+ " disk = get_model_size(model)\n",
+ " params = params_count if params_count is not None else get_num_parameters(model)\n",
+ " return SizeMetrics(disk_bytes=disk, size_mib=_bytes_to_mib(disk), num_params=params)"
+ ]
},
{
"cell_type": "code",
@@ -100,15 +185,61 @@
"id": "52ea5e13",
"metadata": {},
"outputs": [],
- "source": "#| hide\n# Quantization preserves the *number* of params (int8 conv/linear have the same\n# weight count as their fp32 twin) — only the bytes shrink. But quantized weights\n# live in packed params outside `.parameters()`, so the naive sum reports 0.\n# Verify: quantized param count is non-zero and ~equal to the fp32 count.\nimport torch.nn as nn\nfrom torch.ao.quantization import get_default_qconfig, QConfigMapping\nfrom torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\nfrom fasterbench.core import _is_quantized\n\nclass _SizeNet(nn.Module):\n \"Small conv+linear net, statically quantizable via FX.\"\n def __init__(self):\n super().__init__()\n self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n def forward(self, x):\n x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n return self.fc(self.flat(self.pool(x)))\n\n_fp = _SizeNet().eval()\n_x = torch.randn(1, 3, 16, 16)\n_fp_params = get_num_parameters(_fp) # fp32 baseline\n\ntorch.backends.quantized.engine = \"x86\"\n_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n_prep = prepare_fx(_SizeNet().eval(), _qmap, example_inputs=(_x,))\n_prep(_x) # calibrate\n_qm = convert_fx(_prep)\nassert _is_quantized(_qm)\n\n# BEFORE the fix, the naive `.parameters()` sum is 0 for a fully quantized model.\ntest_eq(sum(p.numel() for p in _qm.parameters() if p.requires_grad), 0)\n\n_q_params = get_num_parameters(_qm) # AFTER the fix\nassert _q_params > 0 # zero is gone\nassert abs(_q_params - _fp_params) <= 0.05 * _fp_params # within ~5% of fp32\n# trainable_only must NOT hide quantized weights (they are never trainable)\ntest_eq(get_num_parameters(_qm, trainable_only=False), _q_params)"
+ "source": [
+ "#| hide\n",
+ "# Quantization preserves the *number* of params (int8 conv/linear have the same\n",
+ "# weight count as their fp32 twin) — only the bytes shrink. But quantized weights\n",
+ "# live in packed params outside `.parameters()`, so the naive sum reports 0.\n",
+ "# Verify: quantized param count is non-zero and ~equal to the fp32 count.\n",
+ "import torch.nn as nn\n",
+ "from torch.ao.quantization import get_default_qconfig, QConfigMapping\n",
+ "from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx\n",
+ "from fasterbench.core import _is_quantized\n",
+ "\n",
+ "class _SizeNet(nn.Module):\n",
+ " \"Small conv+linear net, statically quantizable via FX.\"\n",
+ " def __init__(self):\n",
+ " super().__init__()\n",
+ " self.c1, self.c2 = nn.Conv2d(3, 16, 3, padding=1), nn.Conv2d(16, 32, 3, padding=1)\n",
+ " self.pool, self.flat, self.fc = nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(32, 10)\n",
+ " def forward(self, x):\n",
+ " x = torch.relu(self.c1(x)); x = torch.relu(self.c2(x))\n",
+ " return self.fc(self.flat(self.pool(x)))\n",
+ "\n",
+ "_fp = _SizeNet().eval()\n",
+ "_x = torch.randn(1, 3, 16, 16)\n",
+ "_fp_params = get_num_parameters(_fp) # fp32 baseline\n",
+ "\n",
+ "torch.backends.quantized.engine = \"x86\"\n",
+ "_qmap = QConfigMapping().set_global(get_default_qconfig(\"x86\"))\n",
+ "_prep = prepare_fx(_SizeNet().eval(), _qmap, example_inputs=(_x,))\n",
+ "_prep(_x) # calibrate\n",
+ "_qm = convert_fx(_prep)\n",
+ "assert _is_quantized(_qm)\n",
+ "\n",
+ "# BEFORE the fix, the naive `.parameters()` sum is 0 for a fully quantized model.\n",
+ "test_eq(sum(p.numel() for p in _qm.parameters() if p.requires_grad), 0)\n",
+ "\n",
+ "_q_params = get_num_parameters(_qm) # AFTER the fix\n",
+ "assert _q_params > 0 # zero is gone\n",
+ "assert abs(_q_params - _fp_params) <= 0.05 * _fp_params # within ~5% of fp32\n",
+ "# trainable_only must NOT hide quantized weights (they are never trainable)\n",
+ "test_eq(get_num_parameters(_qm, trainable_only=False), _q_params)"
+ ]
},
{
"cell_type": "markdown",
- "execution_count": null,
"id": "82eecedd",
"metadata": {},
- "outputs": [],
- "source": "---\n\n## See Also\n\n- [Benchmark](../analysis/benchmark.html) - Unified benchmarking with `benchmark()`\n- [Compute](compute.html) - MACs and operation counts\n- [Profiling](../analysis/profiling.html) - Per-layer size analysis"
+ "source": [
+ "---\n",
+ "\n",
+ "## See Also\n",
+ "\n",
+ "- [Benchmark](../analysis/benchmark.html) - Unified benchmarking with `benchmark()`\n",
+ "- [Compute](compute.html) - MACs and operation counts\n",
+ "- [Profiling](../analysis/profiling.html) - Per-layer size analysis"
+ ]
}
],
"metadata": {
diff --git a/nbs/metrics/speed.ipynb b/nbs/metrics/speed.ipynb
index 3102546..630896b 100644
--- a/nbs/metrics/speed.ipynb
+++ b/nbs/metrics/speed.ipynb
@@ -34,7 +34,20 @@
"id": "e6b40d9e",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\nimport math, time, warnings\nfrom dataclasses import dataclass, asdict\nfrom typing import Mapping, Sequence\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.benchmark import Timer\n\nfrom fasterbench.core import _device_ctx, _sync, _run_on_devices, _ensure_device_supported"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "import math, time, warnings\n",
+ "from dataclasses import dataclass, asdict\n",
+ "from typing import Mapping, Sequence\n",
+ "\n",
+ "import numpy as np\n",
+ "import torch\n",
+ "import torch.nn as nn\n",
+ "from torch.utils.benchmark import Timer\n",
+ "\n",
+ "from fasterbench.core import _device_ctx, _on_device, _sync, _run_on_devices, _ensure_device_supported"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +55,180 @@
"id": "eb5e3063",
"metadata": {},
"outputs": [],
- "source": "#| export\n@dataclass(slots=True)\nclass SpeedMetrics:\n \"\"\"Latency and throughput metrics for a single device.\"\"\"\n p50_ms: float\n p90_ms: float\n p99_ms: float\n mean_ms: float\n std_ms: float\n throughput_s: float\n\n def as_dict(self) -> dict[str, float]:\n return asdict(self)\n\n\n#| export\ndef _nan_speed_metrics(device: str) -> SpeedMetrics: # device string (unused, for consistent signature)\n \"\"\"Create SpeedMetrics with NaN values for failed benchmarks.\"\"\"\n nan = float(\"nan\")\n return SpeedMetrics(nan, nan, nan, nan, nan, nan)\n\n\n#| export\ndef _stats(lat_ms: np.ndarray, batch: int) -> Mapping[str, float]:\n \"\"\"Compute latency statistics from raw measurements.\"\"\"\n p50, p90, p99 = np.percentile(lat_ms, [50, 90, 99])\n mean = float(lat_ms.mean())\n return {\n \"p50_ms\": float(p50),\n \"p90_ms\": float(p90),\n \"p99_ms\": float(p99),\n \"mean_ms\": mean,\n \"std_ms\": float(lat_ms.std(ddof=1)) if lat_ms.size > 1 else 0.0,\n \"throughput_s\": batch * 1000.0 / mean,\n }\n\n\n#| export\ndef _forward_latencies(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n use_torch_timer: bool | None = None, # force torch.utils.benchmark.Timer (default: auto)\n) -> np.ndarray:\n \"\"\"Collect forward-pass latencies (ms) using optimal timing for each device.\"\"\"\n if use_torch_timer is None:\n use_torch_timer = torch.device(device).type == \"cpu\"\n\n with _device_ctx(device) as dev:\n _ensure_device_supported(model, dev) # quantized ops are CPU-only; avoid SIGSEGV\n model.eval().to(dev)\n sample = sample.to(dev, non_blocking=True)\n\n for _ in range(warmup):\n model(sample)\n _sync(dev)\n\n lat: list[float] = []\n if use_torch_timer and dev.type == \"cpu\":\n t = Timer(stmt=\"model(x)\", globals={\"model\": model, \"x\": sample})\n m = t.blocked_autorange(min_run_time=0.3)\n per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3\n lat = per_iter.tolist()\n elif dev.type == \"cuda\":\n start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)\n for _ in range(steps):\n start_evt.record()\n model(sample)\n end_evt.record()\n _sync(dev)\n lat.append(start_evt.elapsed_time(end_evt))\n else:\n for _ in range(steps):\n t0 = time.perf_counter()\n model(sample)\n lat.append((time.perf_counter() - t0) * 1e3)\n return np.asarray(lat, dtype=np.float32)\n\n\n#| export\n@torch.no_grad()\ndef compute_speed(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n device: str | torch.device = \"cpu\", # device to run on\n warmup: int = 20, # warmup iterations\n steps: int = 100, # measurement iterations\n) -> SpeedMetrics:\n \"\"\"Measure latency and throughput on a single device.\"\"\"\n lat = _forward_latencies(model, sample, device=device, warmup=warmup, steps=steps)\n return SpeedMetrics(**_stats(lat, sample.size(0)))\n\n\n#| export\ndef compute_speed_multi(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n *,\n devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda if available)\n **kwargs,\n) -> dict[str, SpeedMetrics]:\n \"\"\"Measure latency/throughput on multiple devices.\"\"\"\n return _run_on_devices(\n compute_speed, model, sample, devices,\n nan_factory=_nan_speed_metrics,\n metric_name=\"Speed\",\n **kwargs\n )\n\n\n#| export\ndef sweep_threads(\n model: nn.Module, # model to benchmark\n sample: torch.Tensor, # input tensor (with batch dimension)\n thread_counts: Sequence[int] = (1, 2, 4, 8), # thread counts to test\n *,\n warmup: int = 20, # warmup iterations per thread count\n steps: int = 100, # measurement iterations per thread count\n) -> list[dict]:\n \"\"\"Sweep CPU thread counts to find optimal parallelism.\"\"\"\n rows = []\n for n in thread_counts:\n torch.set_num_threads(n)\n lat = _forward_latencies(model, sample, device=\"cpu\", warmup=warmup, steps=steps, use_torch_timer=True)\n rows.append({\"threads\": n, **_stats(lat, sample.size(0))})\n return rows\n\n\n#| export\ndef sweep_batch_sizes(\n model: nn.Module, # model to benchmark\n input_shape: Sequence[int], # input shape WITHOUT batch dim, e.g. (3, 224, 224)\n batch_sizes: Sequence[int] = (1, 2, 4, 8, 16, 32), # batch sizes to test\n *,\n device: str | torch.device = \"cuda\", # device to run on\n warmup: int = 20, # warmup iterations per batch size\n steps: int = 100, # measurement iterations per batch size\n) -> list[dict]:\n \"\"\"Sweep batch sizes to find optimal throughput.\"\"\"\n rows = []\n for bs in batch_sizes:\n try:\n dummy = torch.randn(bs, *input_shape)\n lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n stats = _stats(lat, bs)\n rows.append({\"batch_size\": bs, \"latency_per_sample_ms\": stats[\"mean_ms\"] / bs, **stats})\n except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n warnings.warn(f\"Batch size {bs} failed (likely OOM): {e}\")\n rows.append({\"batch_size\": bs, \"mean_ms\": float(\"nan\"), \"throughput_s\": float(\"nan\")})\n return rows\n\n\n#| export\ndef sweep_latency(\n model: nn.Module, # model to benchmark\n shapes: Sequence[Sequence[int]], # input shapes to test, e.g. [(1,3,224,224), (1,3,384,384)]\n *,\n device: str | torch.device = \"cuda\", # device to run on\n warmup: int = 20, # warmup iterations per shape\n steps: int = 100, # measurement iterations per shape\n) -> list[dict]:\n \"\"\"Sweep input shapes to analyze latency vs resolution.\"\"\"\n rows = []\n for shape in shapes:\n dummy = torch.empty(*shape)\n lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n rows.append({\"shape\": \"×\".join(map(str, shape)), **_stats(lat, shape[0])})\n return rows"
+ "source": [
+ "#| export\n",
+ "@dataclass(slots=True)\n",
+ "class SpeedMetrics:\n",
+ " \"\"\"Latency and throughput metrics for a single device.\"\"\"\n",
+ " p50_ms: float\n",
+ " p90_ms: float\n",
+ " p99_ms: float\n",
+ " mean_ms: float\n",
+ " std_ms: float\n",
+ " throughput_s: float\n",
+ "\n",
+ " def as_dict(self) -> dict[str, float]:\n",
+ " return asdict(self)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _nan_speed_metrics(device: str) -> SpeedMetrics: # device string (unused, for consistent signature)\n",
+ " \"\"\"Create SpeedMetrics with NaN values for failed benchmarks.\"\"\"\n",
+ " nan = float(\"nan\")\n",
+ " return SpeedMetrics(nan, nan, nan, nan, nan, nan)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _stats(lat_ms: np.ndarray, batch: int) -> Mapping[str, float]:\n",
+ " \"\"\"Compute latency statistics from raw measurements.\"\"\"\n",
+ " p50, p90, p99 = np.percentile(lat_ms, [50, 90, 99])\n",
+ " mean = float(lat_ms.mean())\n",
+ " return {\n",
+ " \"p50_ms\": float(p50),\n",
+ " \"p90_ms\": float(p90),\n",
+ " \"p99_ms\": float(p99),\n",
+ " \"mean_ms\": mean,\n",
+ " \"std_ms\": float(lat_ms.std(ddof=1)) if lat_ms.size > 1 else 0.0,\n",
+ " \"throughput_s\": batch * 1000.0 / mean,\n",
+ " }\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def _forward_latencies(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " device: str | torch.device = \"cpu\", # device to run on\n",
+ " warmup: int = 20, # warmup iterations\n",
+ " steps: int = 100, # measurement iterations\n",
+ " use_torch_timer: bool | None = None, # force torch.utils.benchmark.Timer (default: auto)\n",
+ ") -> np.ndarray:\n",
+ " \"\"\"Collect forward-pass latencies (ms) using optimal timing for each device.\"\"\"\n",
+ " if use_torch_timer is None:\n",
+ " use_torch_timer = torch.device(device).type == \"cpu\"\n",
+ "\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",
+ " sample = sample.to(dev, non_blocking=True)\n",
+ "\n",
+ " for _ in range(warmup):\n",
+ " model(sample)\n",
+ " _sync(dev)\n",
+ "\n",
+ " lat: list[float] = []\n",
+ " if use_torch_timer and dev.type == \"cpu\":\n",
+ " t = Timer(stmt=\"model(x)\", globals={\"model\": model, \"x\": sample})\n",
+ " m = t.blocked_autorange(min_run_time=0.3)\n",
+ " per_iter = (np.asarray(m.raw_times) / m.number_per_run) * 1e3\n",
+ " lat = per_iter.tolist()\n",
+ " elif dev.type == \"cuda\":\n",
+ " start_evt, end_evt = torch.cuda.Event(True), torch.cuda.Event(True)\n",
+ " for _ in range(steps):\n",
+ " start_evt.record()\n",
+ " model(sample)\n",
+ " end_evt.record()\n",
+ " _sync(dev)\n",
+ " lat.append(start_evt.elapsed_time(end_evt))\n",
+ " else:\n",
+ " for _ in range(steps):\n",
+ " t0 = time.perf_counter()\n",
+ " model(sample)\n",
+ " lat.append((time.perf_counter() - t0) * 1e3)\n",
+ " return np.asarray(lat, dtype=np.float32)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "@torch.no_grad()\n",
+ "def compute_speed(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " device: str | torch.device = \"cpu\", # device to run on\n",
+ " warmup: int = 20, # warmup iterations\n",
+ " steps: int = 100, # measurement iterations\n",
+ ") -> SpeedMetrics:\n",
+ " \"\"\"Measure latency and throughput on a single device.\"\"\"\n",
+ " lat = _forward_latencies(model, sample, device=device, warmup=warmup, steps=steps)\n",
+ " return SpeedMetrics(**_stats(lat, sample.size(0)))\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def compute_speed_multi(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " *,\n",
+ " devices: Sequence[str | torch.device] | None = None, # devices to benchmark (default: cpu + cuda if available)\n",
+ " **kwargs,\n",
+ ") -> dict[str, SpeedMetrics]:\n",
+ " \"\"\"Measure latency/throughput on multiple devices.\"\"\"\n",
+ " return _run_on_devices(\n",
+ " compute_speed, model, sample, devices,\n",
+ " nan_factory=_nan_speed_metrics,\n",
+ " metric_name=\"Speed\",\n",
+ " **kwargs\n",
+ " )\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def sweep_threads(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " sample: torch.Tensor, # input tensor (with batch dimension)\n",
+ " thread_counts: Sequence[int] = (1, 2, 4, 8), # thread counts to test\n",
+ " *,\n",
+ " warmup: int = 20, # warmup iterations per thread count\n",
+ " steps: int = 100, # measurement iterations per thread count\n",
+ ") -> list[dict]:\n",
+ " \"\"\"Sweep CPU thread counts to find optimal parallelism.\"\"\"\n",
+ " rows = []\n",
+ " for n in thread_counts:\n",
+ " torch.set_num_threads(n)\n",
+ " lat = _forward_latencies(model, sample, device=\"cpu\", warmup=warmup, steps=steps, use_torch_timer=True)\n",
+ " rows.append({\"threads\": n, **_stats(lat, sample.size(0))})\n",
+ " return rows\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def sweep_batch_sizes(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " input_shape: Sequence[int], # input shape WITHOUT batch dim, e.g. (3, 224, 224)\n",
+ " batch_sizes: Sequence[int] = (1, 2, 4, 8, 16, 32), # batch sizes to test\n",
+ " *,\n",
+ " device: str | torch.device = \"cuda\", # device to run on\n",
+ " warmup: int = 20, # warmup iterations per batch size\n",
+ " steps: int = 100, # measurement iterations per batch size\n",
+ ") -> list[dict]:\n",
+ " \"\"\"Sweep batch sizes to find optimal throughput.\"\"\"\n",
+ " rows = []\n",
+ " for bs in batch_sizes:\n",
+ " try:\n",
+ " dummy = torch.randn(bs, *input_shape)\n",
+ " lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n",
+ " stats = _stats(lat, bs)\n",
+ " rows.append({\"batch_size\": bs, \"latency_per_sample_ms\": stats[\"mean_ms\"] / bs, **stats})\n",
+ " except (RuntimeError, torch.cuda.OutOfMemoryError) as e:\n",
+ " warnings.warn(f\"Batch size {bs} failed (likely OOM): {e}\")\n",
+ " rows.append({\"batch_size\": bs, \"mean_ms\": float(\"nan\"), \"throughput_s\": float(\"nan\")})\n",
+ " return rows\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "def sweep_latency(\n",
+ " model: nn.Module, # model to benchmark\n",
+ " shapes: Sequence[Sequence[int]], # input shapes to test, e.g. [(1,3,224,224), (1,3,384,384)]\n",
+ " *,\n",
+ " device: str | torch.device = \"cuda\", # device to run on\n",
+ " warmup: int = 20, # warmup iterations per shape\n",
+ " steps: int = 100, # measurement iterations per shape\n",
+ ") -> list[dict]:\n",
+ " \"\"\"Sweep input shapes to analyze latency vs resolution.\"\"\"\n",
+ " rows = []\n",
+ " for shape in shapes:\n",
+ " dummy = torch.empty(*shape)\n",
+ " lat = _forward_latencies(model, dummy, device=device, warmup=warmup, steps=steps)\n",
+ " rows.append({\"shape\": \"×\".join(map(str, shape)), **_stats(lat, shape[0])})\n",
+ " return rows"
+ ]
},
{
"cell_type": "code",
@@ -81,7 +267,25 @@
"metadata": {},
"outputs": [],
"source": [
- "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_s = compute_speed(_m, _x, warmup=2, steps=5)\nassert isinstance(_s, SpeedMetrics)\nassert _s.mean_ms > 0"
+ "#| hide\nfrom fastcore.test import *\n\nimport torch, torch.nn as nn\n_m = nn.Linear(10, 5)\n_x = torch.randn(1, 10)\n_s = compute_speed(_m, _x, warmup=2, steps=5)\nassert isinstance(_s, SpeedMetrics)\nassert _s.mean_ms > 0\n\n# measuring leaves the model on the device it came in on\ntest_eq({p.device.type for p in _m.parameters()}, {\"cpu\"})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "f8a31c78",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#| hide\n",
+ "#| notest\n",
+ "# Needs a second device: timing a CUDA model on the CPU must not park it there.\n",
+ "if torch.cuda.is_available():\n",
+ " _cm = nn.Linear(10, 5).cuda()\n",
+ " compute_speed(_cm, _x, device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})\n",
+ " sweep_latency(_cm, [(1, 10)], device=\"cpu\", warmup=1, steps=2)\n",
+ " test_eq({p.device.type for p in _cm.parameters()}, {\"cuda\"})"
]
},
{
diff --git a/nbs/tutorials/roofline.ipynb b/nbs/tutorials/roofline.ipynb
index 340a42a..8ebd2c0 100644
--- a/nbs/tutorials/roofline.ipynb
+++ b/nbs/tutorials/roofline.ipynb
@@ -13,7 +13,7 @@
"id": "callout",
"metadata": {},
"source": [
- "> This notebook demonstrates measurement primitives. For compression decisions based on roofline data, see fasterrecipes."
+ "> This notebook demonstrates measurement primitives. This library measures; the decisions that use these numbers live in the compression workflow that calls it."
]
},
{
@@ -75,7 +75,13 @@
"cell_type": "markdown",
"id": "resnet_header",
"metadata": {},
- "source": "## Profiling ResNet-18\n\n`RooflineAnalyzer` profiles a model in a single pass under the hood: forward hooks on every leaf module measure FLOPs (computed analytically for Conv and Linear), bytes moved (weights + input + output per Williams 2009), and wall time.\n\nIf you do not pass a `peaks=` argument, it calls `measure_peaks()` automatically."
+ "source": [
+ "## Profiling ResNet-18\n",
+ "\n",
+ "`RooflineAnalyzer` profiles a model in a single pass under the hood: forward hooks on every leaf module measure FLOPs (computed analytically for Conv and Linear), bytes moved (weights + input + output per Williams 2009), and wall time.\n",
+ "\n",
+ "If you do not pass a `peaks=` argument, it calls `measure_peaks()` automatically."
+ ]
},
{
"cell_type": "code",
@@ -99,7 +105,19 @@
"cell_type": "markdown",
"id": "resnet_example_out",
"metadata": {},
- "source": "```\n=== Roofline =============================================================\n name type FLOPs bytes AI GFLOPs/s bound util%\n layer4.0.conv2 Conv2d 231.21M 10.01M 23.10 820.14 compute 10.1%\n layer4.1.conv1 Conv2d 231.21M 10.01M 23.10 810.22 compute 10.0%\n layer4.1.conv2 Conv2d 231.21M 10.01M 23.10 812.49 compute 10.0%\n layer3.0.conv2 Conv2d 115.61M 5.11M 22.62 402.33 compute 5.0%\n ...\n```\n\nEach row shows a layer's FLOPs, bytes moved, arithmetic intensity, achieved throughput, bound classification, and utilization (fraction of the roof reached)."
+ "source": [
+ "```\n",
+ "=== Roofline =============================================================\n",
+ " name type FLOPs bytes AI GFLOPs/s bound util%\n",
+ " layer4.0.conv2 Conv2d 231.21M 10.01M 23.10 820.14 compute 10.1%\n",
+ " layer4.1.conv1 Conv2d 231.21M 10.01M 23.10 810.22 compute 10.0%\n",
+ " layer4.1.conv2 Conv2d 231.21M 10.01M 23.10 812.49 compute 10.0%\n",
+ " layer3.0.conv2 Conv2d 115.61M 5.11M 22.62 402.33 compute 5.0%\n",
+ " ...\n",
+ "```\n",
+ "\n",
+ "Each row shows a layer's FLOPs, bytes moved, arithmetic intensity, achieved throughput, bound classification, and utilization (fraction of the roof reached)."
+ ]
},
{
"cell_type": "markdown",
@@ -179,7 +197,20 @@
"cell_type": "markdown",
"id": "summary",
"metadata": {},
- "source": "## Summary\n\n| Tool | Purpose |\n|------|---------|\n| `measure_peaks()` | Empirically probe peak FLOPs/s and streaming bandwidth |\n| `HardwarePeaks` | Dataclass holding device peaks and ridge point |\n| `RooflineAnalyzer` | Per-layer roofline profiler |\n| `RooflineAnalyzer.profile()` | Measure FLOPs, bytes moved, and time per layer |\n| `RooflineAnalyzer.summary()` | Print a table of the slowest layers with their roofline metrics |\n| `RooflineAnalyzer.plot()` | Plotly figure with roof ceiling and per-layer markers |\n| `RooflinePoint` | Dataclass for a single layer's measurement |\n| `clear_peaks_cache()` | Reset the `measure_peaks()` cache |"
+ "source": [
+ "## Summary\n",
+ "\n",
+ "| Tool | Purpose |\n",
+ "|------|---------|\n",
+ "| `measure_peaks()` | Empirically probe peak FLOPs/s and streaming bandwidth |\n",
+ "| `HardwarePeaks` | Dataclass holding device peaks and ridge point |\n",
+ "| `RooflineAnalyzer` | Per-layer roofline profiler |\n",
+ "| `RooflineAnalyzer.profile()` | Measure FLOPs, bytes moved, and time per layer |\n",
+ "| `RooflineAnalyzer.summary()` | Print a table of the slowest layers with their roofline metrics |\n",
+ "| `RooflineAnalyzer.plot()` | Plotly figure with roof ceiling and per-layer markers |\n",
+ "| `RooflinePoint` | Dataclass for a single layer's measurement |\n",
+ "| `clear_peaks_cache()` | Reset the `measure_peaks()` cache |"
+ ]
},
{
"cell_type": "markdown",
diff --git a/nbs/visualization/plot.ipynb b/nbs/visualization/plot.ipynb
index 1ac297a..2103d38 100644
--- a/nbs/visualization/plot.ipynb
+++ b/nbs/visualization/plot.ipynb
@@ -34,7 +34,15 @@
"id": "0c86ca8f-38d2-44d0-8204-3c977f1b1c19",
"metadata": {},
"outputs": [],
- "source": "#| export\nfrom __future__ import annotations\n\nfrom fasterbench.utils import *\nimport re, numbers\nfrom typing import Any, Mapping, Sequence\nimport plotly.graph_objects as go"
+ "source": [
+ "#| export\n",
+ "from __future__ import annotations\n",
+ "\n",
+ "from fasterbench.utils import *\n",
+ "import re, numbers\n",
+ "from typing import Any, Mapping, Sequence\n",
+ "import plotly.graph_objects as go"
+ ]
},
{
"cell_type": "code",
@@ -42,7 +50,107 @@
"id": "7309e3bc",
"metadata": {},
"outputs": [],
- "source": "#| export\n_UNIT_RX = re.compile(r\"([.\\d]+)\\s*([kKmMgGtT]?)[bB]?$\")\n\n#| export\n_SCALE: dict[str, float] = dict(K=1e3, M=1e6, G=1e9, T=1e12,\n k=1e3, m=1e6, g=1e9, t=1e12)\n\n#| export\ndef _parse(x): # metric value (number or string with unit suffix)\n \"\"\"Parse a metric value that may include unit suffixes (K, M, G, T).\"\"\"\n if isinstance(x, numbers.Number): return float(x)\n m = _UNIT_RX.match(str(x))\n if not m: raise ValueError(f\"Cannot parse metric value: {x}\")\n num, suf = m.groups()\n return float(num) / _SCALE.get(suf, 1.0)\n\n\n#| export\n# Default specs: (extractor, format_string, default_max)\nSPECS: dict[str, tuple] = {\n \"💾\": (lambda d: d[\"size_disk_bytes\"] / 1e6, \"Size : {:.1f} MiB\", 500),\n \"🧮\": (lambda d: _parse(d[\"size_num_params\"]) / 1e6, \"Params: {:.2f} M\", 50),\n \"⏱️\": (lambda d: d[\"speed_cpu_mean_ms\"], \"Latency: {:.1f} ms\", 2000),\n \"🌍\": (lambda d: _parse(d[\"energy_cpu_co2_eq_g\"]) * 1e3, \"Emission: {:.3f} mgCO2e\", 5),\n \"🔋\": (lambda d: d[\"energy_cpu_energy_wh\"] * 1e3, \"Energy: {:.3f} mWh\", 10),\n}\n\n#| export\ndef create_radar_plot(\n benchmark_results: Sequence[Mapping[str, Any]], # sequence of benchmark result dicts\n model_names: Sequence[str] | None = None, # optional list of model names for the legend\n reference_max: Mapping[str, float] | None = None, # optional dict mapping emoji to max value\n specs: dict[str, tuple] | None = None, # custom specs override (defaults to SPECS)\n) -> go.Figure:\n \"\"\"Create a radar plot comparing benchmark results across models.\"\"\"\n specs = specs or SPECS\n\n if model_names is None:\n model_names = [f\"Model {i}\" for i in range(len(benchmark_results))]\n\n # ---------- collect raw values -------------------------------------\n raw_matrix: dict[str, list[float]] = {e: [] for e in specs}\n hover_mat: dict[str, list[str]] = {e: [] for e in specs}\n\n for res in benchmark_results:\n for emoji, (extract, fmt, _) in specs.items():\n val = extract(res)\n raw_matrix[emoji].append(val)\n hover_mat[emoji].append(fmt.format(val))\n\n # ---------- reference range (min=0) --------------------------------\n ref_max = {e: max(max(vals), default) # default from specs\n for e, (_, _, default) in specs.items()\n for vals in [raw_matrix[e]]}\n if reference_max:\n ref_max.update(reference_max)\n\n # ---------- normalise & build traces -------------------------------\n labels = list(specs.keys()) + [list(specs.keys())[0]] # close loop\n fig = go.Figure()\n\n for idx, name in enumerate(model_names):\n norm: list[float] = []\n hovers: list[str] = []\n for e in specs:\n val = raw_matrix[e][idx]\n norm.append(val / ref_max[e] if ref_max[e] else 0)\n hovers.append(hover_mat[e][idx])\n norm += norm[:1]\n hovers += hovers[:1]\n\n fig.add_trace(\n go.Scatterpolar(\n r=norm,\n theta=labels,\n name=name,\n hovertext=hovers,\n hoverinfo=\"text\",\n fill=\"toself\",\n )\n )\n\n # ---------- layout --------------------------------------------------\n fig.update_layout(\n polar=dict(\n radialaxis=dict(range=[0, 1], visible=True,\n showticklabels=False,\n gridcolor=\"rgba(128,128,128,0.4)\",\n linecolor=\"rgba(128,128,128,0.4)\"),\n angularaxis=dict(tickfont=dict(size=24),\n gridcolor=\"rgba(128,128,128,0.4)\"),\n bgcolor=\"rgba(0,0,0,0)\",\n ),\n paper_bgcolor=\"rgba(0,0,0,0)\",\n plot_bgcolor=\"rgba(0,0,0,0)\",\n margin=dict(t=60, b=60, l=60, r=60),\n showlegend=len(model_names) > 1,\n )\n return fig"
+ "source": [
+ "#| export\n",
+ "_UNIT_RX = re.compile(r\"([.\\d]+)\\s*([kKmMgGtT]?)[bB]?$\")\n",
+ "\n",
+ "#| export\n",
+ "_SCALE: dict[str, float] = dict(K=1e3, M=1e6, G=1e9, T=1e12,\n",
+ " k=1e3, m=1e6, g=1e9, t=1e12)\n",
+ "\n",
+ "#| export\n",
+ "def _parse(x): # metric value (number or string with unit suffix)\n",
+ " \"\"\"Parse a metric value that may include unit suffixes (K, M, G, T).\"\"\"\n",
+ " if isinstance(x, numbers.Number): return float(x)\n",
+ " m = _UNIT_RX.match(str(x))\n",
+ " if not m: raise ValueError(f\"Cannot parse metric value: {x}\")\n",
+ " num, suf = m.groups()\n",
+ " return float(num) / _SCALE.get(suf, 1.0)\n",
+ "\n",
+ "\n",
+ "#| export\n",
+ "# Default specs: (extractor, format_string, default_max)\n",
+ "SPECS: dict[str, tuple] = {\n",
+ " \"💾\": (lambda d: d[\"size_disk_bytes\"] / 1e6, \"Size : {:.1f} MiB\", 500),\n",
+ " \"🧮\": (lambda d: _parse(d[\"size_num_params\"]) / 1e6, \"Params: {:.2f} M\", 50),\n",
+ " \"⏱️\": (lambda d: d[\"speed_cpu_mean_ms\"], \"Latency: {:.1f} ms\", 2000),\n",
+ " \"🌍\": (lambda d: _parse(d[\"energy_cpu_co2_eq_g\"]) * 1e3, \"Emission: {:.3f} mgCO2e\", 5),\n",
+ " \"🔋\": (lambda d: d[\"energy_cpu_energy_wh\"] * 1e3, \"Energy: {:.3f} mWh\", 10),\n",
+ "}\n",
+ "\n",
+ "#| export\n",
+ "def create_radar_plot(\n",
+ " benchmark_results: Sequence[Mapping[str, Any]], # sequence of benchmark result dicts\n",
+ " model_names: Sequence[str] | None = None, # optional list of model names for the legend\n",
+ " reference_max: Mapping[str, float] | None = None, # optional dict mapping emoji to max value\n",
+ " specs: dict[str, tuple] | None = None, # custom specs override (defaults to SPECS)\n",
+ ") -> go.Figure:\n",
+ " \"\"\"Create a radar plot comparing benchmark results across models.\"\"\"\n",
+ " specs = specs or SPECS\n",
+ "\n",
+ " if model_names is None:\n",
+ " model_names = [f\"Model {i}\" for i in range(len(benchmark_results))]\n",
+ "\n",
+ " # ---------- collect raw values -------------------------------------\n",
+ " raw_matrix: dict[str, list[float]] = {e: [] for e in specs}\n",
+ " hover_mat: dict[str, list[str]] = {e: [] for e in specs}\n",
+ "\n",
+ " for res in benchmark_results:\n",
+ " for emoji, (extract, fmt, _) in specs.items():\n",
+ " val = extract(res)\n",
+ " raw_matrix[emoji].append(val)\n",
+ " hover_mat[emoji].append(fmt.format(val))\n",
+ "\n",
+ " # ---------- reference range (min=0) --------------------------------\n",
+ " ref_max = {e: max(max(vals), default) # default from specs\n",
+ " for e, (_, _, default) in specs.items()\n",
+ " for vals in [raw_matrix[e]]}\n",
+ " if reference_max:\n",
+ " ref_max.update(reference_max)\n",
+ "\n",
+ " # ---------- normalise & build traces -------------------------------\n",
+ " labels = list(specs.keys()) + [list(specs.keys())[0]] # close loop\n",
+ " fig = go.Figure()\n",
+ "\n",
+ " for idx, name in enumerate(model_names):\n",
+ " norm: list[float] = []\n",
+ " hovers: list[str] = []\n",
+ " for e in specs:\n",
+ " val = raw_matrix[e][idx]\n",
+ " norm.append(val / ref_max[e] if ref_max[e] else 0)\n",
+ " hovers.append(hover_mat[e][idx])\n",
+ " norm += norm[:1]\n",
+ " hovers += hovers[:1]\n",
+ "\n",
+ " fig.add_trace(\n",
+ " go.Scatterpolar(\n",
+ " r=norm,\n",
+ " theta=labels,\n",
+ " name=name,\n",
+ " hovertext=hovers,\n",
+ " hoverinfo=\"text\",\n",
+ " fill=\"toself\",\n",
+ " )\n",
+ " )\n",
+ "\n",
+ " # ---------- layout --------------------------------------------------\n",
+ " fig.update_layout(\n",
+ " polar=dict(\n",
+ " radialaxis=dict(range=[0, 1], visible=True,\n",
+ " showticklabels=False,\n",
+ " gridcolor=\"rgba(128,128,128,0.4)\",\n",
+ " linecolor=\"rgba(128,128,128,0.4)\"),\n",
+ " angularaxis=dict(tickfont=dict(size=24),\n",
+ " gridcolor=\"rgba(128,128,128,0.4)\"),\n",
+ " bgcolor=\"rgba(0,0,0,0)\",\n",
+ " ),\n",
+ " paper_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " plot_bgcolor=\"rgba(0,0,0,0)\",\n",
+ " margin=dict(t=60, b=60, l=60, r=60),\n",
+ " showlegend=len(model_names) > 1,\n",
+ " )\n",
+ " return fig"
+ ]
},
{
"cell_type": "code",