fix(metrics): correct params & MACs for quantized models

fasterbench/_modidx.py

@@ -35,6 +35,8 @@
                                                                                      'fasterbench/compute.py'),
                                      'fasterbench.compute.ComputeMetrics.macs_available': ( 'metrics/compute.html#computemetrics.macs_available',
                                                                                             'fasterbench/compute.py'),
+                                     'fasterbench.compute._count_macs_hooks': ( 'metrics/compute.html#_count_macs_hooks',
+                                                                                'fasterbench/compute.py'),
                                      'fasterbench.compute.compute_compute': ( 'metrics/compute.html#compute_compute',
                                                                               'fasterbench/compute.py')},
             'fasterbench.core': { 'fasterbench.core._bytes_to_mib': ('core/core.html#_bytes_to_mib', 'fasterbench/core.py'),
@@ -166,6 +168,8 @@
                                                                               'fasterbench/roofline.py')},
             'fasterbench.size': { 'fasterbench.size.SizeMetrics': ('metrics/size.html#sizemetrics', 'fasterbench/size.py'),
                                   'fasterbench.size.SizeMetrics.as_dict': ('metrics/size.html#sizemetrics.as_dict', 'fasterbench/size.py'),
+                                  'fasterbench.size._quantized_param_count': ( 'metrics/size.html#_quantized_param_count',
+                                                                               'fasterbench/size.py'),
                                   'fasterbench.size.compute_size': ('metrics/size.html#compute_size', 'fasterbench/size.py'),
                                   'fasterbench.size.get_model_size': ('metrics/size.html#get_model_size', 'fasterbench/size.py'),
                                   'fasterbench.size.get_num_parameters': ('metrics/size.html#get_num_parameters', 'fasterbench/size.py')},

fasterbench/compute.py

@@ -3,12 +3,15 @@
 # %% ../nbs/metrics/compute.ipynb #0091d170
 from __future__ import annotations
 
+import math
 import warnings
 from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
 
+from .core import _is_quantized
+
 try:
     from thop import profile as _thop_profile
 except ImportError:
@@ -39,6 +42,48 @@ def as_dict(self) -> dict[str, float]:
         }
 
 
+#| export
+def _count_macs_hooks(
+    model: nn.Module,      # model to profile (run on CPU)
+    sample: torch.Tensor,  # input tensor (with batch dimension)
+) -> int:
+    """Count Conv2d/Linear MACs via forward hooks (fp *and* quantized modules).
+
+    thop has no handlers for quantized ops (it reports ~0), so we count manually.
+    Quantization does not change the number of multiply-accumulates: an int8 conv
+    does the same MACs as its fp32 twin. Uses module attributes (`out_channels`,
+    `kernel_size`, `in_features`, ...) which both fp and quantized modules expose,
+    and reads output spatial dims from the produced tensor.
+    """
+    macs = 0
+    handles = []
+
+    def conv_hook(m, inp, out):
+        nonlocal macs
+        out_spatial = math.prod(out.shape[2:])                       # out_H * out_W
+        kernel = math.prod(m.kernel_size)                            # kH * kW
+        macs += out.shape[0] * m.out_channels * (m.in_channels // m.groups) * kernel * out_spatial
+
+    def linear_hook(m, inp, out):
+        nonlocal macs
+        leading = out.numel() // m.out_features                      # batch/token dims
+        macs += leading * m.in_features * m.out_features
+
+    for m in model.modules():
+        if hasattr(m, "kernel_size") and hasattr(m, "out_channels") and hasattr(m, "groups"):
+            handles.append(m.register_forward_hook(conv_hook))
+        elif hasattr(m, "in_features") and hasattr(m, "out_features"):
+            handles.append(m.register_forward_hook(linear_hook))
+
+    try:
+        with torch.no_grad():
+            model(sample)
+    finally:
+        for h in handles:
+            h.remove()
+    return macs
+
+
 #| export
 def compute_compute(
     model: nn.Module,        # model to analyze
@@ -55,6 +100,15 @@ def compute_compute(
 
     macs_m: float | None = None
 
+    # Quantized models: thop has no quantized-op handlers (returns ~0). Count the
+    # Conv2d/Linear MACs manually — quantization preserves the MAC count.
+    if _is_quantized(model):
+        try:
+            macs_m = round(_count_macs_hooks(model, sample) / 1e6, 3)
+        except Exception as e:
+            warnings.warn(f"quantized MAC counting failed: {e}")
+        return ComputeMetrics(macs_m=macs_m)
+
     if _thop_profile is not None:
         try:
             mac_raw, _ = _thop_profile(model, inputs=(sample,), verbose=False)

fasterbench/size.py

@@ -3,7 +3,7 @@
 # %% ../nbs/metrics/size.ipynb #4c37777e
 from __future__ import annotations
 
-from .core import _bytes_to_mib
+from .core import _bytes_to_mib, _is_quantized
 import torch
 import io
 from dataclasses import dataclass, asdict
@@ -22,15 +22,47 @@ def get_model_size(model: torch.nn.Module) -> int:  # model to measure
     return buf.getbuffer().nbytes
 
 
+#| export
+def _quantized_param_count(model: torch.nn.Module) -> int:  # model with quantized modules
+    """Sum the numel of packed quantized weights/biases (not in `.parameters()`).
+
+    Quantized `Conv2d`/`Linear` keep their weights in `_packed_params`, exposed via
+    a *callable* `.weight()`/`.bias()` accessor (unlike a plain `nn.Parameter`, which
+    is a tensor). Quantization preserves the number of parameters, so counting these
+    makes a quantized model's param count match its fp32 twin.
+    """
+    total = 0
+    for m in model.modules():
+        for name in ("weight", "bias"):
+            accessor = getattr(m, name, None)
+            if not callable(accessor):  # plain nn.Parameter tensors are not callable
+                continue
+            try:
+                t = accessor()
+                if t is not None:
+                    total += t.numel()
+            except Exception:
+                pass
+    return total
+
+
 #| export
 def get_num_parameters(
     model: torch.nn.Module,       # model to count parameters
     trainable_only: bool = True,  # if True, only count trainable parameters
 ) -> int:
-    """Count the number of (optionally trainable) parameters."""
+    """Count the number of (optionally trainable) parameters.
+
+    Quantized weights live in packed params outside `.parameters()`, so they are
+    added separately (regardless of `trainable_only`, as they are never trainable).
+    """
     if trainable_only:
-        return sum(p.numel() for p in model.parameters() if p.requires_grad)
-    return sum(p.numel() for p in model.parameters())
+        n = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    else:
+        n = sum(p.numel() for p in model.parameters())
+    if _is_quantized(model):
+        n += _quantized_param_count(model)
+    return n
 
 
 #| export

nbs/metrics/compute.ipynb

@@ -34,34 +34,15 @@
    "id": "0091d170",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "#| export\n",
-    "from __future__ import annotations\n",
-    "\n",
-    "import warnings\n",
-    "from dataclasses import dataclass\n",
-    "\n",
-    "import torch\n",
-    "import torch.nn as nn\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"
-   ]
+   "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"
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "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 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    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)\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)"
   },
   {
    "cell_type": "code",
@@ -101,6 +82,14 @@
     "#| 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_c = compute_compute(_m, _x)\nassert isinstance(_c, ComputeMetrics)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "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"
+  },
   {
    "cell_type": "markdown",
    "id": "see_also",