espdlx 0.1.0__py3-none-any.whl

espdlx/__init__.py

@@ -0,0 +1,7 @@
+"""espdlx: esp-dl friendly NN blocks for ESP32-S3 (PyTorch backend)."""
+
+__version__ = "0.1.0"
+
+from .graph import Model
+
+__all__ = ["Model", "__version__"]

espdlx/convert.py

@@ -0,0 +1,314 @@
+"""ONNX export + esp-dl quantization for espdlx models.
+
+This is the shipped version of the pipeline previously duplicated across the
+experiment scripts (``flower_to_espdl.py``, ``stl10_to_espdl.py``,
+``micro_to_espdl.py`` — not shipped)::
+
+    from espdlx.convert import convert
+
+    report = convert(model, example_input, calib_loader, "model.espdl")
+
+Steps: ``torch.onnx.export`` -> esp-dl friendly rewrite
+(``Relu + Min(6)`` -> ``Clip``, ``Reshape`` -> ``Flatten``) ->
+``espdl_quantize_onnx`` (ESP-PPQ) -> ``.espdl`` + report.
+
+Requires the ``convert`` extra (``pip install espdlx[convert]``).
+``onnx`` / ``esp_ppq`` are imported lazily so ``import espdlx`` stays light.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+__all__ = [
+    "convert",
+    "export_onnx",
+    "make_espdl_friendly",
+    "quantize_onnx",
+]
+
+DEFAULT_TARGET = "esp32s3"
+DEFAULT_QUANT_TYPE = "w8a8"
+DEFAULT_OPSET = 13
+DEFAULT_CALIB_STEPS = 16
+
+
+def _require_onnx() -> Any:
+    try:
+        import onnx
+    except ImportError as exc:
+        raise ImportError(
+            "espdlx.convert needs the 'convert' extra "
+            "(pip install espdlx[convert]): missing package 'onnx'"
+        ) from exc
+    return onnx
+
+
+def _require_espdl_quantize() -> Any:
+    try:
+        from esp_ppq.api import espdl_quantize_onnx
+    except ImportError as exc:
+        raise ImportError(
+            "espdlx.convert needs the 'convert' extra "
+            "(pip install espdlx[convert]): missing package 'esp-ppq'"
+        ) from exc
+    return espdl_quantize_onnx
+
+
+def export_onnx(
+    model: Any,
+    example_input: Any,
+    onnx_path: str | Path,
+    *,
+    input_name: str = "input",
+    output_name: str = "output",
+    opset: int = DEFAULT_OPSET,
+) -> Path:
+    """Export ``model`` to ONNX (classic exporter, fixed opset).
+
+    Returns the resolved ``onnx_path``. The model is left untouched
+    (callers typically ``model.eval()`` first — see :func:`convert`).
+    """
+    import torch
+
+    onnx_path = Path(onnx_path)
+    onnx_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.onnx.export(
+        model,
+        example_input,
+        str(onnx_path),
+        input_names=[input_name],
+        output_names=[output_name],
+        opset_version=opset,
+        do_constant_folding=True,
+    )
+    return onnx_path
+
+
+def make_espdl_friendly(model: Any) -> Any:
+    """Rewrite ONNX exporter artifacts to ops esp-dl supports, in place.
+
+    - ``Relu + Min(6)`` -> ``Clip(0, 6)`` (esp-dl has Clip, no Min)
+    - ``Reshape`` (static) -> ``Flatten`` (old ppq executor chokes on Reshape)
+
+    Takes and returns the ``onnx.ModelProto``. Runs ``onnx.checker`` first
+    (import) and last (result) so failures surface here, not in quantization.
+    """
+    onnx = _require_onnx()
+    onnx.checker.check_model(model)
+
+    inits = {i.name: onnx.numpy_helper.to_array(i) for i in model.graph.initializer}
+    six = [n for n, v in inits.items() if v.shape == () and float(v) == 6.0]
+    if six:
+        six_name = six[0]
+    else:
+        import numpy as _np
+
+        six_name = "clip_max_6"
+        model.graph.initializer.append(
+            onnx.numpy_helper.from_array(_np.array(6.0, dtype=_np.float32), six_name)
+        )
+    zero_name = "clip_min_0"
+    if zero_name not in inits:
+        import numpy as _np
+
+        model.graph.initializer.append(
+            onnx.numpy_helper.from_array(_np.array(0.0, dtype=_np.float32), zero_name)
+        )
+
+    by_out = {}
+    for n in model.graph.node:
+        for o in n.output:
+            by_out[o] = n
+    relu_outs = set()
+    for n in model.graph.node:
+        if n.op_type == "Min" and six_name in n.input:
+            key = n.input[0] if n.input[1] == six_name else n.input[1]
+            relu = by_out.get(key, None)
+            if relu is not None and relu.op_type == "Relu" and key == relu.output[0]:
+                relu_outs.add(key)
+
+    new_nodes = []
+    for n in model.graph.node:
+        if n.op_type == "Relu" and n.output[0] in relu_outs:
+            continue  # fused into the Clip below
+        if n.op_type == "Min" and six_name in n.input:
+            key = n.input[0] if n.input[1] == six_name else n.input[1]
+            relu = by_out.get(key, None)
+            if relu is not None and key in relu_outs:
+                new_nodes.append(
+                    onnx.helper.make_node(
+                        "Clip",
+                        inputs=[relu.input[0], zero_name, six_name],
+                        outputs=[n.output[0]],
+                        name=(n.name + "_clip") if n.name else "",
+                    )
+                )
+                continue
+        if n.op_type == "Reshape":
+            flat = onnx.helper.make_node(
+                "Flatten",
+                inputs=[n.input[0]],
+                outputs=[n.output[0]],
+                name=(n.name + "_flat") if n.name else "",
+            )
+            flat.attribute.append(onnx.helper.make_attribute("axis", 1))
+            new_nodes.append(flat)
+            continue
+        new_nodes.append(n)
+    del model.graph.node[:]
+    model.graph.node.extend(new_nodes)
+
+    used = {i for n in model.graph.node for i in n.input}
+    for i in list(model.graph.initializer):
+        if i.name == "val_1" and i.name not in used:
+            model.graph.initializer.remove(i)
+    onnx.checker.check_model(model)
+    return model
+
+
+def _default_collate(batch: Any) -> Any:
+    """Default PPQ collate: ``DataLoader(TensorDataset(x))`` yields ``[x]``."""
+    if isinstance(batch, (list, tuple)) and len(batch) == 1:
+        batch = batch[0]
+    to_cpu = getattr(batch, "to", None)
+    return to_cpu("cpu") if callable(to_cpu) else batch
+
+
+def quantize_onnx(
+    onnx_path: str | Path,
+    espdl_path: str | Path,
+    calib_loader: Any,
+    input_shape: list[int] | tuple[int, ...],
+    *,
+    target: str = DEFAULT_TARGET,
+    quant_type: str = DEFAULT_QUANT_TYPE,
+    calib_steps: int = DEFAULT_CALIB_STEPS,
+    collate_fn: Any = None,
+    device: str = "cpu",
+    error_report: bool = False,
+    export_test_values: bool = False,
+    verbose: int = 0,
+) -> Any:
+    """Quantize a (friendly) ONNX model via ESP-PPQ; returns the PPQ graph."""
+    espdl_quantize_onnx = _require_espdl_quantize()
+    espdl_path = Path(espdl_path)
+    espdl_path.parent.mkdir(parents=True, exist_ok=True)
+    return espdl_quantize_onnx(
+        onnx_import_file=str(onnx_path),
+        espdl_export_file=str(espdl_path),
+        calib_dataloader=calib_loader,
+        calib_steps=calib_steps,
+        input_shape=list(input_shape),
+        target=target,
+        quant_type=quant_type,
+        collate_fn=collate_fn or _default_collate,
+        device=device,
+        error_report=error_report,
+        export_test_values=export_test_values,
+        verbose=verbose,
+    )
+
+
+def _input_quant_params(graph: Any, input_name: str) -> tuple[Any, Any]:
+    """Best-effort input scale / zero-point from a quantized PPQ graph."""
+    try:
+        g_in = graph.inputs[input_name]
+        cfg = g_in.dest_op_configs[0]
+        sc, off = cfg.scale, cfg.offset
+        import torch
+
+        sc = float(sc.reshape(-1)[0]) if torch.is_tensor(sc) else float(sc)
+        off = int(off.reshape(-1)[0]) if torch.is_tensor(off) else int(off)
+        return sc, off
+    except Exception:
+        return None, None
+
+
+def convert(
+    model: Any,
+    example_input: Any,
+    calib_loader: Any,
+    espdl_path: str | Path,
+    *,
+    onnx_path: str | Path | None = None,
+    report_path: str | Path | None = "auto",
+    input_shape: list[int] | tuple[int, ...] | None = None,
+    input_name: str = "input",
+    output_name: str = "output",
+    opset: int = DEFAULT_OPSET,
+    target: str = DEFAULT_TARGET,
+    quant_type: str = DEFAULT_QUANT_TYPE,
+    calib_steps: int = DEFAULT_CALIB_STEPS,
+    collate_fn: Any = None,
+    device: str = "cpu",
+    verbose: int = 0,
+) -> dict:
+    """Export ``model`` and quantize to ``.espdl``; return a JSON-able report.
+
+    - ``example_input``: dummy tensor for ``torch.onnx.export``
+      (also derives ``input_shape`` when omitted).
+    - ``calib_loader``: calibration batches (e.g. ``DataLoader(TensorDataset)``).
+    - ``espdl_path``: where to write the ``.espdl`` file.
+    - ``onnx_path``: defaults to ``espdl_path`` with a ``.onnx`` suffix.
+    - ``report_path``: ``"auto"`` writes ``espdl_report.json`` next to the
+      ``.espdl`` file (scratch convention); pass a path to override,
+      ``None`` to skip writing (the report dict is still returned).
+    """
+    onnx = _require_onnx()
+    espdl_path = Path(espdl_path)
+    if onnx_path is None:
+        onnx_path = espdl_path.with_suffix(".onnx")
+    else:
+        onnx_path = Path(onnx_path)
+    if input_shape is None:
+        input_shape = [int(d) for d in example_input.shape]
+
+    model.eval()
+    export_onnx(
+        model,
+        example_input,
+        onnx_path,
+        input_name=input_name,
+        output_name=output_name,
+        opset=opset,
+    )
+    friendly = onnx.load(str(onnx_path))
+    make_espdl_friendly(friendly)
+    onnx.save(friendly, str(onnx_path))
+
+    graph = quantize_onnx(
+        onnx_path,
+        espdl_path,
+        calib_loader,
+        input_shape,
+        target=target,
+        quant_type=quant_type,
+        calib_steps=calib_steps,
+        collate_fn=collate_fn,
+        device=device,
+        verbose=verbose,
+    )
+    scale, zero_point = _input_quant_params(graph, input_name)
+    report = {
+        "onnx_path": str(onnx_path),
+        "espdl_path": str(espdl_path),
+        "espdl_bytes": espdl_path.stat().st_size,
+        "opset": opset,
+        "target": target,
+        "quant_type": quant_type,
+        "calib_steps": calib_steps,
+        "input_name": input_name,
+        "input_shape": list(input_shape),
+        "input_scale": scale,
+        "input_zero_point": zero_point,
+    }
+    if report_path == "auto":
+        report_path = espdl_path.parent / "espdl_report.json"
+    if report_path is not None:
+        report_path = Path(report_path)
+        report_path.parent.mkdir(parents=True, exist_ok=True)
+        report_path.write_text(json.dumps(report, indent=2))
+    return report

espdlx/fomo.py

@@ -0,0 +1,180 @@
+"""EXPERIMENTAL FOMO helpers: the two things Edge Impulse does better.
+
+Status: synthetic unit tests only (`tests/test_fomo.py`); best real-data
+run so far reached P=0.03/R=0.28 — not a working detector yet.
+
+Kept deliberately small (no parity chase):
+
+- logits-only training (softmax at export/decode, never in the model),
+- imbalance handling: object-weighted cross-entropy + classifier bias init
+  from the dataset foreground prior.
+
+Targets are per-cell class maps ``(N, G, G)`` with ``0`` = background,
+``1`` = object centroid cell (see :func:`encode_centroids`).
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+__all__ = [
+    "encode_centroids",
+    "encode_heatmap",
+    "weighted_loss",
+    "soft_loss",
+    "init_bias_from_prior",
+    "decode",
+    "decode_peaks",
+]
+
+
+def encode_centroids(
+    boxes: list[tuple[float, float, float, float]],
+    grid: int,
+) -> torch.Tensor:
+    """Map normalized ``(cx, cy, w, h)`` boxes to a ``(G, G)`` class map.
+
+    The cell containing each centroid is marked ``1``; everything else stays
+    ``0``. Zero/negative-area boxes are skipped.
+    """
+    target = torch.zeros(grid, grid, dtype=torch.long)
+    for cx, cy, w, h in boxes:
+        if w <= 0 or h <= 0:
+            continue
+        j = min(max(int(math.floor(cx * grid)), 0), grid - 1)
+        i = min(max(int(math.floor(cy * grid)), 0), grid - 1)
+        target[i, j] = 1
+    return target
+
+
+def weighted_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    object_weight: float = 100.0,
+) -> torch.Tensor:
+    """Weighted cross-entropy over ``(N, 2, G, G)`` logits.
+
+    Background weight is ``1.0``; the object cell weight defaults to ``100.0``
+    (Edge Impulse default) to counter the ~1:143 foreground/background ratio
+    on a 12x12 grid.
+    """
+    weight = torch.tensor([1.0, float(object_weight)], device=logits.device)
+    return F.cross_entropy(logits, targets.long(), weight=weight)
+
+
+def encode_heatmap(
+    boxes: list[tuple[float, float, float, float]],
+    grid: int,
+    sigma: float = 1.0,
+    radius: int = 2,
+) -> torch.Tensor:
+    """Splat normalized ``(cx, cy, w, h)`` boxes into a ``(G, G)`` heatmap.
+
+    Each centroid stamps a Gaussian bump (peak 1 at the exact centroid,
+    ``exp(-d^2 / 2*sigma^2)`` in cell units, ``radius`` cells wide);
+    overlapping bumps merge by max. Zero/negative-area boxes are skipped.
+    Pairs with :func:`soft_loss` / :func:`decode_peaks`.
+    """
+    heat = torch.zeros(grid, grid)
+    for cx, cy, w, h in boxes:
+        if w <= 0 or h <= 0:
+            continue
+        gx, gy = cx * grid, cy * grid
+        jc, ic = int(math.floor(gx)), int(math.floor(gy))
+        for i in range(max(ic - radius, 0), min(ic + radius, grid - 1) + 1):
+            for j in range(max(jc - radius, 0), min(jc + radius, grid - 1) + 1):
+                d2 = (gx - j) ** 2 + (gy - i) ** 2
+                v = math.exp(-d2 / (2 * sigma * sigma))
+                if v > heat[i, j]:
+                    heat[i, j] = v
+    return heat
+
+
+def soft_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    heatmaps: torch.Tensor,
+    object_weight: float = 20.0,
+    gamma: float = 4.0,
+) -> torch.Tensor:
+    """Localization-tolerant loss over ``(N, 2, G, G)`` logits.
+
+    - positives (centroid cells from :func:`encode_centroids`): NLL toward
+      the object logit, scaled by ``object_weight``;
+    - negatives: NLL toward background, scaled by ``(1 - heat)^gamma`` with
+      ``heat`` from :func:`encode_heatmap` — an adjacent cell (heat ~0.6)
+      pays ~2% of a far cell's penalty, so "roughly right" barely hurts.
+    Normalized by the positive-cell count (min 1), so the scale is stable
+    across images with different object counts.
+    """
+    logp = F.log_softmax(logits, dim=1)
+    pos = targets.long() == 1
+    n_pos = max(int(pos.sum().item()), 1)
+    loss_pos = -logp[:, 1][pos].sum() * float(object_weight)
+    w = (1.0 - heatmaps.float()).pow(float(gamma))
+    loss_neg = -(w * logp[:, 0])[~pos].sum()
+    return (loss_pos + loss_neg) / n_pos
+
+
+def init_bias_from_prior(model: torch.nn.Module, obj_prior: float) -> None:
+    """Init the final head bias so the model starts at the dataset prior.
+
+    With a 2-logit softmax head, ``bias = [0, log(p / (1 - p))]`` makes the
+    initial object probability equal the foreground cell fraction ``p``
+    (Edge Impulse's ``set_classifier_biases_from_dataset`` idea, reduced to
+    one number). Finds the last ``espdlx`` conv layer in ``model.layers``.
+    """
+    from .layers.conv import Conv2d as EspConv2d
+
+    p = min(max(float(obj_prior), 1e-6), 1.0 - 1e-6)
+    head = None
+    for lyr in getattr(model, "layers", []):
+        if isinstance(lyr, EspConv2d) and lyr.out_channels == 2:
+            head = lyr
+    if head is None:
+        raise ValueError("init_bias_from_prior: no 2-channel head conv found")
+    with torch.no_grad():
+        bias = head.conv.bias
+        assert bias is not None and bias.numel() == 2
+        bias.zero_()
+        bias[1] = math.log(p / (1.0 - p))
+
+
+@torch.no_grad()
+def decode(
+    logits: torch.Tensor,
+    thresh: float = 0.5,
+) -> list[list[tuple[int, int, float]]]:
+    """Logits ``(N, 2, G, G)`` -> per image ``[(col, row, score)]`` centroids.
+
+    Softmax is applied here (or at export), never during training. ``score``
+    is the object-channel probability.
+    """
+    prob = torch.softmax(logits, dim=1)[:, 1]
+    out: list[list[tuple[int, int, float]]] = []
+    for n in range(logits.shape[0]):
+        idx = (prob[n] > thresh).nonzero(as_tuple=False)
+        out.append([(int(j), int(i), float(prob[n, i, j])) for i, j in idx.tolist()])
+    return out
+
+
+@torch.no_grad()
+def decode_peaks(
+    logits: torch.Tensor,
+    thresh: float = 0.5,
+) -> list[list[tuple[int, int, float]]]:
+    """Peak-NMS decode: 3x3 max-pool keeps local maxima above ``thresh``.
+
+    Pairs with :func:`soft_loss` — the loss lets bumps spread over adjacent
+    cells, this collapses each bump back to one centroid.
+    """
+    prob = torch.softmax(logits, dim=1)[:, 1]
+    peaks = F.max_pool2d(prob[:, None], 3, stride=1, padding=1)[:, 0] == prob
+    out: list[list[tuple[int, int, float]]] = []
+    for n in range(logits.shape[0]):
+        idx = ((prob[n] > thresh) & peaks[n]).nonzero(as_tuple=False)
+        out.append([(int(j), int(i), float(prob[n, i, j])) for i, j in idx.tolist()])
+    return out

espdlx/graph.py

@@ -0,0 +1,37 @@
+"""The `Model` container: an ordered stack of espdlx layers.
+
+Deploy with :func:`espdlx.convert.convert` (``torch.onnx.export`` ->
+ESP-PPQ ``espdl_quantize_onnx`` -> ``.espdl``).
+"""
+
+from __future__ import annotations
+
+from torch import nn
+
+from .layers.base import Layer
+
+
+class Model(nn.Module):
+    def __init__(self, layers: list, name: str = "model") -> None:
+        super().__init__()
+        for lyr in layers:
+            if not isinstance(lyr, Layer):
+                raise TypeError(
+                    f"{type(lyr).__name__} is not an espdlx layer; use espdlx.layers"
+                )
+        self.name = name
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, x):
+        for lyr in self.layers:
+            x = lyr(x)
+        return x
+
+    def validate_shapes(self, input_shape: tuple) -> tuple:
+        shape = tuple(input_shape)
+        for lyr in self.layers:
+            shape = lyr.validate_shapes(shape)
+        return shape
+
+    def __repr__(self) -> str:
+        return f"Model(name={self.name!r}, layers={len(self.layers)})"

espdlx/layers/__init__.py

@@ -0,0 +1,27 @@
+"""espdlx constrained layer blocks (torch nn.Module wrappers)."""
+
+from .activations import BatchNorm2d, HardSwish, ReLU, ReLU6, Sigmoid, Softmax
+from .base import Layer
+from .conv import Conv2d, DepthwiseConv2d
+from .linear import Linear
+from .math import Add, Flatten, Mean, Mul
+from .pool import AvgPool2d, MaxPool2d
+
+__all__ = [
+    "Layer",
+    "Conv2d",
+    "DepthwiseConv2d",
+    "Linear",
+    "MaxPool2d",
+    "AvgPool2d",
+    "ReLU",
+    "ReLU6",
+    "HardSwish",
+    "Sigmoid",
+    "Softmax",
+    "BatchNorm2d",
+    "Add",
+    "Mul",
+    "Mean",
+    "Flatten",
+]

espdlx/layers/activations.py

@@ -0,0 +1,70 @@
+"""Activation blocks (plain torch semantics, esp-dl deployable)."""
+
+from __future__ import annotations
+
+import torch.nn.functional as F
+from torch import nn
+
+from .base import Layer
+
+
+class ReLU(Layer):
+    def forward(self, x):
+        return F.relu(x)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+class ReLU6(Layer):
+    def forward(self, x):
+        return F.relu6(x)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+class HardSwish(Layer):
+    def forward(self, x):
+        return F.hardswish(x)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+class Sigmoid(Layer):
+    def forward(self, x):
+        return F.sigmoid(x)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+class Softmax(Layer):
+    def __init__(self, dim: int = -1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        return F.softmax(x, dim=self.dim)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+class BatchNorm2d(Layer):
+    """Batch normalization (standard torch semantics, exports to ONNX)."""
+
+    def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
+        super().__init__()
+        self.num_features = int(num_features)
+        self.bn = nn.BatchNorm2d(self.num_features, eps=eps, momentum=momentum)
+
+    def forward(self, x):
+        return self.bn(x)
+
+    def validate_shapes(self, in_shape):
+        return in_shape
+
+
+__all__ = ["ReLU", "ReLU6", "HardSwish", "Sigmoid", "Softmax", "BatchNorm2d"]

espdlx/layers/base.py

@@ -0,0 +1,11 @@
+"""Shared base for espdlx constrained blocks."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class Layer(nn.Module):
+    def validate_shapes(self, in_shape: tuple) -> tuple:
+        raise NotImplementedError