rpx_benchmark.model_profiler

  1"""Unified model profiler for all RPX tasks.
  2
  3One class, any model, all metrics. No per-task boilerplate.
  4
  5Usage
  6-----
  7    from rpx_benchmark.model_profiler import ModelProfiler
  8
  9    profiler = ModelProfiler(model)
 10    report = profiler.profile()
 11
 12    # Hardware-agnostic (same on any machine)
 13    print(report.params_m)              # 524.4
 14    print(report.flops_g)               # 187.3
 15
 16    # Measured (hardware-specific, tagged with system card)
 17    print(report.latency_p50_ms)        # 142.7
 18    print(report.memory_peak_mb)        # 3200.0
 19
 20    # Deployability verdicts (novel — directly answers "can I use this?")
 21    print(report.realtime_at_hz(10))    # False (needs 142ms, budget is 100ms)
 22    print(report.realtime_at_hz(5))     # True  (needs 142ms, budget is 200ms)
 23    print(report.budget_fraction(10))   # 1.42  (exceeds 10Hz budget by 42%)
 24
 25    # Get the EfficiencyMetadata for the runner
 26    eff = report.to_efficiency_metadata()
 27
 28Design
 29------
 30The profiler handles two distinct phases:
 31
 32**Pre-run** (called once, before inference):
 33    Params, FLOPs, memory traffic estimate, system card.
 34    These are model-intrinsic or estimated — no data needed.
 35
 36**Post-run** (filled by the runner during inference):
 37    Latency percentiles, peak memory, FLOP counter refinement.
 38    The profiler returns an :class:`EfficiencyMetadata` that the runner
 39    mutates in-place during ``run_with_deployment_readiness``.
 40
 41This split means the profiler works for any task — depth, pose, segmentation,
 42detection — without knowing what task it is. The task script just does::
 43
 44    profiler = ModelProfiler(model_or_adapter)
 45    eff = profiler.pre_run_profile()
 46    # ... pass eff to runner, runner fills latency/memory ...
 47    report = profiler.full_report(eff)
 48
 49Novel metric: Real-Time Budget Analysis
 50----------------------------------------
 51For robotics deployment, the key question isn't "what's the latency?" but
 52"does it fit in my control loop?"  RPX reports a **budget fraction**:
 53
 54.. math::
 55
 56    \\beta(f) = \\frac{\\tau_{\\text{model}}}{1/f}
 57
 58where :math:`f` is the target control frequency (Hz) and
 59:math:`\\tau_{\\text{model}}` is measured inference latency.
 60
 61:math:`\\beta < 1` means the model fits within the real-time budget.
 62:math:`\\beta > 1` means it exceeds it, and by how much.
 63
 64Reported at standard robot control frequencies: 5 Hz (slow manipulation),
 6510 Hz (fast manipulation), 30 Hz (navigation).  No existing perception
 66benchmark reports this.
 67"""
 68
 69from __future__ import annotations
 70
 71from dataclasses import dataclass
 72from typing import Any, Dict, Optional, Tuple
 73
 74from .logging_utils import get_logger
 75from .profiler import (
 76    EfficiencyMetadata,
 77    SystemCard,
 78    count_parameters,
 79    estimate_memory_traffic_gb,
 80)
 81
 82log = get_logger(__name__)
 83
 84# Standard robot control frequencies for budget analysis.
 85ROBOT_CONTROL_HZ: Tuple[float, ...] = (5.0, 10.0, 30.0)
 86"""Slow manipulation (5 Hz), fast manipulation (10 Hz), navigation (30 Hz)."""
 87
 88
 89def _find_torch_module(model: Any) -> Any:
 90    """Walk common attribute paths to find an ``nn.Module`` with ``.parameters()``.
 91
 92    Works for BenchmarkableModel, raw adapters, and torch modules directly.
 93    Returns ``None`` if no torch module is found.
 94    """
 95    # If it already has .parameters(), it IS a torch module
 96    if hasattr(model, "parameters") and callable(model.parameters):
 97        return model
 98
 99    # Walk common adapter layouts
100    for path in (
101        "model",
102        "torch_module",
103        "_model",
104        "_pipe.model",
105        "adapter.model",
106        "adapter._model",
107        "adapter.torch_module",
108    ):
109        cur = model
110        for part in path.split("."):
111            cur = getattr(cur, part, None)
112            if cur is None:
113                break
114        if cur is not None and hasattr(cur, "parameters") and callable(cur.parameters):
115            return cur
116
117    return None
118
119
120@dataclass
121class ProfileReport:
122    """Complete profiling result for a model.
123
124    Tier 1 and budget analysis are the primary outputs.
125    Tier 3 fields are filled after the runner completes.
126    """
127
128    # ── Tier 1: Hardware-agnostic ────────────────────────────────────────
129    params_m: Optional[float] = None
130    flops_g: Optional[float] = None
131    macs_g: Optional[float] = None
132    memory_traffic_gb: Optional[float] = None
133    arithmetic_intensity: Optional[float] = None
134
135    # ── Tier 2: Roofline (hardware-parametric) ───────────────────────────
136    roofline: Optional[Dict[str, Dict[str, float]]] = None
137
138    # ── Tier 3: Measured (hardware-specific) ─────────────────────────────
139    latency_p50_ms: Optional[float] = None
140    latency_p95_ms: Optional[float] = None
141    latency_p99_ms: Optional[float] = None
142    peak_cpu_mb: Optional[float] = None
143    peak_cuda_mb: Optional[float] = None
144    peak_mps_mb: Optional[float] = None
145    system_card: Optional[Dict[str, Any]] = None
146
147    # ── Budget analysis (novel) ──────────────────────────────────────────
148    budget_fractions: Optional[Dict[str, float]] = None
149    """Maps ``"5hz"`` → β, ``"10hz"`` → β, ``"30hz"`` → β."""
150
151    def budget_fraction(self, hz: float) -> Optional[float]:
152        """Compute β = latency / (1000/hz).  < 1 means real-time feasible."""
153        if self.latency_p50_ms is None or hz <= 0:
154            return None
155        budget_ms = 1000.0 / hz
156        return round(self.latency_p50_ms / budget_ms, 3)
157
158    def realtime_at_hz(self, hz: float) -> Optional[bool]:
159        """Can this model sustain ``hz`` frames per second?"""
160        beta = self.budget_fraction(hz)
161        if beta is None:
162            return None
163        return beta <= 1.0
164
165    def compute_budget_fractions(self) -> Dict[str, float]:
166        """Compute budget fractions at standard robot control frequencies."""
167        out: Dict[str, float] = {}
168        for hz in ROBOT_CONTROL_HZ:
169            beta = self.budget_fraction(hz)
170            if beta is not None:
171                out[f"{hz:.0f}hz"] = beta
172        self.budget_fractions = out
173        return out
174
175    def to_dict(self) -> Dict[str, Any]:
176        """Structured output for JSON serialization."""
177        return {
178            "tier1_hardware_agnostic": {
179                "params_m": self.params_m,
180                "flops_g": self.flops_g,
181                "macs_g": self.macs_g,
182                "memory_traffic_gb": self.memory_traffic_gb,
183                "arithmetic_intensity": self.arithmetic_intensity,
184            },
185            "tier2_roofline": self.roofline,
186            "tier3_measured": {
187                "latency_p50_ms": self.latency_p50_ms,
188                "latency_p95_ms": self.latency_p95_ms,
189                "latency_p99_ms": self.latency_p99_ms,
190                "peak_cpu_mb": self.peak_cpu_mb,
191                "peak_cuda_mb": self.peak_cuda_mb,
192                "peak_mps_mb": self.peak_mps_mb,
193                "system_card": self.system_card,
194            },
195            "realtime_budget": self.budget_fractions,
196        }
197
198    def summary_line(self) -> str:
199        """One-line summary for terminal output."""
200        parts = []
201        if self.params_m is not None:
202            parts.append(f"{self.params_m:.1f}M params")
203        if self.flops_g is not None:
204            parts.append(f"{self.flops_g:.1f}G FLOPs")
205        if self.latency_p50_ms is not None:
206            parts.append(f"{self.latency_p50_ms:.1f}ms p50")
207        if self.budget_fractions:
208            for hz_key, beta in self.budget_fractions.items():
209                symbol = "✓" if beta <= 1.0 else "✗"
210                parts.append(f"{symbol}@{hz_key}")
211        return " | ".join(parts) if parts else "no profile data"
212
213
214class ModelProfiler:
215    """Unified profiler for any RPX model.  Works across all tasks.
216
217    Parameters
218    ----------
219    model : Any
220        A BenchmarkableModel, adapter, torch nn.Module, or callable.
221        The profiler auto-discovers the torch module for Tier 1 metrics.
222    input_shape : tuple
223        ``(C, H, W)`` for FLOPs estimation.  Default: RPX standard 640×480 RGB.
224    device : str
225        Device for profiling.  Default ``"cuda"``.
226
227    Example
228    -------
229    ::
230
231        profiler = ModelProfiler(my_adapter)
232
233        # Pre-run: get Tier 1 metrics + EfficiencyMetadata for the runner
234        eff = profiler.pre_run_profile()
235
236        # ... runner.run_with_deployment_readiness(efficiency=eff) ...
237        # (runner fills latency, memory, FLOPs into eff)
238
239        # Post-run: build the full report from the now-populated eff
240        report = profiler.full_report(eff)
241        print(report.summary_line())
242        # → "524.4M params | 187.3G FLOPs | 142.7ms p50 | ✗@10hz | ✓@5hz"
243    """
244
245    def __init__(
246        self,
247        model: Any,
248        input_shape: Tuple[int, ...] = (3, 480, 640),
249        device: str = "cuda",
250    ) -> None:
251        self.model = model
252        self.input_shape = input_shape
253        self.device = device
254        self._torch_mod = _find_torch_module(model)
255
256    @property
257    def has_torch_module(self) -> bool:
258        return self._torch_mod is not None
259
260    def pre_run_profile(self) -> EfficiencyMetadata:
261        """Compute Tier 1 metrics and return an EfficiencyMetadata for the runner.
262
263        The returned object is passed to
264        ``runner.run_with_deployment_readiness(efficiency=eff)``.
265        The runner will fill in Tier 3 fields (latency, memory, FLOPs)
266        during the benchmark run.
267        """
268        eff = EfficiencyMetadata(
269            model_type="local" if self._torch_mod is not None else "classical",
270        )
271
272        if self._torch_mod is not None:
273            try:
274                eff.params_m = count_parameters(self._torch_mod)
275            except Exception as e:
276                log.warning("params count failed: %s", e)
277
278            try:
279                eff.memory_traffic_gb = estimate_memory_traffic_gb(
280                    self._torch_mod, self.input_shape, device=self.device,
281                )
282            except Exception as e:
283                log.warning("memory traffic estimate failed: %s", e)
284
285        try:
286            eff.system_card = SystemCard.auto_detect(
287                input_resolution=f"{self.input_shape[-1]}x{self.input_shape[-2]}",
288            )
289        except Exception as e:
290            log.warning("system card detection failed: %s", e)
291
292        return eff
293
294    def full_report(self, eff: EfficiencyMetadata) -> ProfileReport:
295        """Build a complete ProfileReport from a (now runner-populated) EfficiencyMetadata.
296
297        Call this AFTER the runner has completed — ``eff`` will have
298        latency, memory, and FLOPs filled in by the runner.
299        """
300        # Ensure derived fields are populated
301        eff.derive_tier1()
302        if eff.flops_g is not None and eff.roofline is None:
303            eff.compute_roofline()
304
305        # Build roofline dict for serialization
306        roofline_dict = None
307        if eff.roofline:
308            roofline_dict = {
309                name: {
310                    "compute_ms": b.compute_ms,
311                    "memory_ms": b.memory_ms,
312                    "latency_ms": b.latency_ms,
313                    "bottleneck": b.bottleneck,
314                }
315                for name, b in eff.roofline.items()
316            }
317
318        report = ProfileReport(
319            # Tier 1
320            params_m=eff.params_m,
321            flops_g=eff.flops_g,
322            macs_g=eff.macs_g,
323            memory_traffic_gb=eff.memory_traffic_gb,
324            arithmetic_intensity=eff.arithmetic_intensity,
325            # Tier 2
326            roofline=roofline_dict,
327            # Tier 3
328            latency_p50_ms=eff.latency_p50_ms,
329            latency_p95_ms=eff.latency_p95_ms,
330            latency_p99_ms=eff.latency_p99_ms,
331            peak_cpu_mb=eff.peak_cpu_mb,
332            peak_cuda_mb=eff.peak_cuda_mb,
333            peak_mps_mb=eff.peak_mps_mb,
334            system_card=eff.system_card.to_dict() if eff.system_card else None,
335        )
336
337        # Compute real-time budget analysis
338        report.compute_budget_fractions()
339
340        return report