API Reference

Auto-generated from the docstrings. See Quick Start for how the pieces fit together.

Layers

The sparse layers live in sparsepixels.layers. InputReduce produces the sparse representation, and the rest consume it.

InputReduce

InputReduce(n=30, threshold=0.0, beta_n=0.005, beta_maskedE=1.0, learn_n=True, learn_threshold=True, tau_threshold=0.05, tau_n=1.0, **kwargs)

Bases: Layer

Reduce a dense image to its first n active pixels for sparse FPGA inference.

Keeps the first n pixels whose first channel is above threshold, in raster order, and zeroes the rest, returning the masked image together with a 0/1 keep mask that the following sparse layers use as the sparse representation.

The budget n and the threshold can be learned during training (the default) so they need not be tuned by hand; set learn_n or learn_threshold to False to keep either fixed (both False gives the plain, non-learnable selection). The selection is always exact -- the learnable versions only shape the gradient, so the layer behaves identically at inference and stays deployable. When n is learned, a penalty of weight beta_n nudges it smaller, trading a little accuracy for lower FPGA latency and resources; it starts at n and can range over [1, H*W]. An optional penalty of weight beta_maskedE discourages masking pixel intensity, giving the threshold (and n) a restoring force so a trainable threshold settles near the noise floor instead of over-masking signal. After training, read the values to deploy from the n_max_pixels and threshold properties.

Parameters:

Name	Type	Description	Default
n		initial pixel budget, and the fixed budget when learn_n is False.	30
threshold		initial activity threshold on the first channel, fixed when learn_threshold is False.	0.0
beta_n		weight of the budget penalty added to the loss (0 disables it).	0.005
beta_maskedE		weight of the masked-intensity penalty (0 disables it). Penalizes the fraction of pixel intensity that gets masked, so masking bright (signal) pixels is costly while masking dim (noise) pixels is cheap -- an adaptive, non-vanishing restoring force against over-masking.	1.0
learn_n		make the pixel budget trainable.	True
learn_threshold		make the threshold trainable.	True
tau_threshold		softness of the threshold surrogate used to obtain gradients.	0.05
tau_n		softness of the budget-cutoff surrogate used to obtain gradients.	1.0

Source code in sparsepixels/layers.py

def __init__(
    self,
    n=30,
    threshold=0.0,
    beta_n=5e-3,
    beta_maskedE=1.0,
    learn_n=True,
    learn_threshold=True,
    tau_threshold=0.05,
    tau_n=1.0,
    **kwargs,
):
    super().__init__(**kwargs)
    self.n_init = int(n)
    self.threshold_init = float(threshold)
    self.beta_n = float(beta_n)
    self.beta_maskedE = float(beta_maskedE)
    self.learn_n = learn_n
    self.learn_threshold = learn_threshold
    self.tau_threshold = float(tau_threshold)
    self.tau_n = float(tau_n)

n_max_pixels property

n_max_pixels

Integer pixel budget to deploy (the initial value until the layer is built).

threshold property

threshold

Threshold to deploy (the initial value until the layer is built).

QConv2DSparse

QConv2DSparse(*conv_args, **conv_kwargs)

Bases: Layer

Quantized 2D convolution that operates on the sparse (active-pixel) representation.

Wraps an HGQ QConv2D: masks the input to the active pixels, convolves, adds a separately quantized per-filter bias on the nonzero outputs, applies the activation, then re-masks the output. This is numerically the same as a dense quantized conv restricted to the active pixels, which is what the HLS sparse_conv kernel computes.

Parameters:

Name	Type	Description	Default
*conv_args		positional arguments forwarded to hgq.layers.QConv2D (e.g. filters, kernel_size).	()
**conv_kwargs		keyword arguments forwarded to QConv2D (padding, strides, ...). use_bias, activation and bq_conf are handled here: the bias has its own weight and quantizer (bq_conf), and the activation is applied after the bias.	{}

Call args

inputs: tuple (x, keep_mask) of the feature map and its keep mask.

Source code in sparsepixels/layers.py

def __init__(self, *conv_args, **conv_kwargs):
    super().__init__(name=conv_kwargs.get("name", None))
    self._use_bias = conv_kwargs.pop("use_bias", True)
    self._bq_conf = conv_kwargs.pop("bq_conf", None) or QuantizerConfig("default", "bias")
    self._activation = keras.activations.get(conv_kwargs.pop("activation", None))

    conv_kwargs["use_bias"] = False
    conv_kwargs["activation"] = None
    self.conv = QConv2D(*conv_args, **conv_kwargs)
    self.masker = RemoveDilatedPixels()

AveragePooling2DSparse

AveragePooling2DSparse(*pool_args, **pool_kwargs)

Bases: Layer

Average pooling on the sparse representation.

Average-pools the feature map and max-pools the keep mask, so a pooled cell stays active when any of its source pixels were active. Mirrors the HLS sparse_pooling_avg kernel.

Parameters:

Name	Type	Description	Default
*pool_args		positional arguments forwarded to keras AveragePooling2D (e.g. pool_size).	()
**pool_kwargs		keyword arguments forwarded to the pooling layers.	{}

Call args

inputs: tuple (x, keep_mask) of the feature map and its keep mask.

Source code in sparsepixels/layers.py

def __init__(self, *pool_args, **pool_kwargs):
    super().__init__(name=pool_kwargs.get("name", None))
    self.avg_pool = AveragePooling2D(*pool_args, **pool_kwargs)
    self.max_pool = MaxPooling2D(*pool_args, **pool_kwargs)

MaxPooling2DSparse

MaxPooling2DSparse(*pool_args, **pool_kwargs)

Bases: Layer

Max pooling on the sparse representation.

Max-pools both the feature map and the keep mask. Mirrors the HLS sparse_pooling_max kernel.

Parameters:

Name	Type	Description	Default
*pool_args		positional arguments forwarded to keras MaxPooling2D (e.g. pool_size).	()
**pool_kwargs		keyword arguments forwarded to the pooling layer.	{}

Call args

inputs: tuple (x, keep_mask) of the feature map and its keep mask.

Source code in sparsepixels/layers.py

def __init__(self, *pool_args, **pool_kwargs):
    super().__init__(name=pool_kwargs.get("name", None))
    self.max_pool = MaxPooling2D(*pool_args, **pool_kwargs)

Data study

Utilities in sparsepixels.utils for picking a threshold and budget before training.

active_pixels_vs_threshold

active_pixels_vs_threshold(x, thresholds=None, channel=0, percentiles=(25, 50, 75), plot=True, ax=None, save_path=None)

Study how the number of active pixels per image changes with the threshold.

For each candidate threshold, counts the active pixels (x[..., channel] > threshold) in every image and reports the mean, min, max and the given percentiles across the dataset. Use it to pick a threshold and an initial budget n: at your chosen threshold, a high percentile or the max tells you how many pixels the busiest images have.

Parameters:

Name	Type	Description	Default
x		image array of shape (N, H, W, C).	required
thresholds		thresholds to scan; defaults to 50 points across the data range.	None
channel		channel used to decide activity (default 0, matching InputReduce).	0
percentiles		percentiles of the per-image active-pixel count to report.	(25, 50, 75)
plot		draw the curves (otherwise only the stats are computed).	True
ax		optional Axes to draw into; when given, the figure is not shown here.	None

Returns a dict with a threshold array plus mean, min, max and p arrays.

Source code in sparsepixels/utils.py

def active_pixels_vs_threshold(
    x, thresholds=None, channel=0, percentiles=(25, 50, 75), plot=True, ax=None, save_path=None
):
    """Study how the number of active pixels per image changes with the threshold.

    For each candidate threshold, counts the active pixels (x[..., channel] > threshold) in every
    image and reports the mean, min, max and the given percentiles across the dataset. Use it to pick
    a threshold and an initial budget n: at your chosen threshold, a high percentile or the max tells
    you how many pixels the busiest images have.

    Args:
        x: image array of shape (N, H, W, C).
        thresholds: thresholds to scan; defaults to 50 points across the data range.
        channel: channel used to decide activity (default 0, matching InputReduce).
        percentiles: percentiles of the per-image active-pixel count to report.
        plot: draw the curves (otherwise only the stats are computed).
        ax: optional Axes to draw into; when given, the figure is not shown here.

    Returns a dict with a threshold array plus mean, min, max and p<k> arrays.
    """
    x = np.asarray(x)
    flat = x[..., channel].reshape(x.shape[0], -1)
    if thresholds is None:
        thresholds = np.linspace(float(flat.min()), float(flat.max()), 50)
    thresholds = np.asarray(thresholds, dtype=float)

    stats = {"threshold": thresholds}
    for key in ("mean", "min", "max", *[f"p{p}" for p in percentiles]):
        stats[key] = np.empty_like(thresholds)
    for i, thr in enumerate(thresholds):
        counts = (flat > thr).sum(axis=1)
        stats["mean"][i] = counts.mean()
        stats["min"][i] = counts.min()
        stats["max"][i] = counts.max()
        for p in percentiles:
            stats[f"p{p}"][i] = np.percentile(counts, p)

    if plot:
        owns_fig = ax is None
        if ax is None:
            fig, ax = plt.subplots(figsize=(7, 4.5))
        ax.plot(thresholds, stats["mean"], color="black", lw=2, label="mean")
        ax.plot(thresholds, stats["min"], "--", alpha=0.7, label="min")
        for p in percentiles:
            ax.plot(thresholds, stats[f"p{p}"], "--", alpha=0.7, label=f"{p}th")
        ax.plot(thresholds, stats["max"], "--", alpha=0.7, label="max")
        ax.set_xlabel("threshold")
        ax.set_ylabel("active pixels per image")
        ax.set_title("active pixels vs threshold")
        ax.legend(loc="best")
        ax.grid(alpha=0.2)
        if owns_fig:
            if save_path is not None:
                fig.savefig(save_path, dpi=150, bbox_inches="tight")
            plt.show()
    return stats

plot_reduced_examples

plot_reduced_examples(x, n, threshold, indices=None, n_examples=5, channel=0, figsize=None, aspect=None, cmap='viridis', save_path=None)

Show what InputReduce keeps for a given budget n and threshold on a few images.

Each row shows the original image and, at the same intensity scale, the pixels that are kept (the first n active pixels in raster order), annotated with the kept and active counts. Use it to check that a candidate n and threshold keep the informative pixels.

Parameters:

Name	Type	Description	Default
x		image array of shape (N, H, W, C).	required
n		pixel budget (the first n active pixels are kept).	required
threshold		activity threshold on the given channel.	required
indices		which images to show; defaults to the first n_examples.	None
n_examples		number of images to show when indices is not given.	5
channel		channel used to decide activity (default 0).	0
figsize		optional (width, height).	None
aspect		imshow aspect: None keeps the natural pixel ratio, "auto" stretches each image to fill its panel (useful for very wide images), or a float scales the pixel height.	None
cmap		colormap for the active pixels; empty pixels (at or below 0) are drawn white. For binary images a reversed map like "viridis_r" makes the hits dark on white.	'viridis'

Source code in sparsepixels/utils.py

def plot_reduced_examples(
    x, n, threshold, indices=None, n_examples=5, channel=0, figsize=None, aspect=None, cmap="viridis", save_path=None
):
    """Show what InputReduce keeps for a given budget n and threshold on a few images.

    Each row shows the original image and, at the same intensity scale, the pixels that are kept (the
    first n active pixels in raster order), annotated with the kept and active counts. Use it to check
    that a candidate n and threshold keep the informative pixels.

    Args:
        x: image array of shape (N, H, W, C).
        n: pixel budget (the first n active pixels are kept).
        threshold: activity threshold on the given channel.
        indices: which images to show; defaults to the first n_examples.
        n_examples: number of images to show when indices is not given.
        channel: channel used to decide activity (default 0).
        figsize: optional (width, height).
        aspect: imshow aspect: None keeps the natural pixel ratio, "auto" stretches each image to
            fill its panel (useful for very wide images), or a float scales the pixel height.
        cmap: colormap for the active pixels; empty pixels (at or below 0) are drawn white. For
            binary images a reversed map like "viridis_r" makes the hits dark on white.
    """
    x = np.asarray(x)
    if indices is None:
        indices = np.arange(min(n_examples, len(x)))
    indices = list(indices)

    cmap_obj = plt.get_cmap(cmap).copy()
    cmap_obj.set_bad("white")

    fig, axes = plt.subplots(len(indices), 2, figsize=figsize or (6, 3 * len(indices)), squeeze=False)
    for row, idx in enumerate(indices):
        s = x[idx][..., channel]
        active = s > threshold
        rank = np.cumsum(active.reshape(-1))
        keep = (active.reshape(-1) & (rank <= n)).reshape(s.shape)
        n_active, n_kept = int(active.sum()), int(keep.sum())
        vmax = float(s.max()) or 1.0

        panels = [
            (s, f"image #{idx}"),
            (np.where(keep, s, 0.0), f"kept (n={n}, thr={threshold:g}): {n_kept} of {n_active} active"),
        ]
        for a, (img, title) in zip(axes[row], panels):
            a.imshow(
                np.ma.masked_less_equal(img, 0.0),
                cmap=cmap_obj,
                vmin=0,
                vmax=vmax,
                aspect=aspect if aspect is not None else "equal",
                interpolation="nearest",
            )
            a.set_title(title)
            a.set_xticks([])
            a.set_yticks([])
    fig.tight_layout()
    if save_path is not None:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.show()

Training & monitoring

SparseTrainingMonitor

Bases: Callback

Record sparse-model diagnostics into the training history each epoch.

Add it to the callbacks in model.fit; it stores these in the History so plot_history can show them next to the loss and your compiled metrics:

• n_max_pixels and threshold: the learned or fixed pixel budget and threshold of the InputReduce layer (suffixed with the layer name when there is more than one).
• ebops: total EBOPS over the quantized layers, a proxy for the quantized hardware cost. This monitor applies set_sparse_ebops_factor at the start of training, so the EBOPS (and its regularizer) reflect the sparse (active-pixel) compute without any manual call.
• loss_task, loss_ebops, loss_n, loss_maskedE: the training loss split into the task loss, the EBOPS penalty, the pixel-budget penalty and the masked-intensity penalty (the last present only when its weight beta_maskedE is nonzero), each already scaled by its weight so the parts add up to the loss.

plot_history

plot_history(history, early_stopping=None, figsize=None, ncols=3, save_path=None)

Plot the training history recorded by SparseTrainingMonitor in one figure.

Shows the total loss and each compiled metric (train, and validation when present), the loss breakdown (task, EBOPS and budget penalties, in red) and the sparse diagnostics (n_max_pixels, threshold, ebops, in green). Panels are only created for keys that are present, so it works with or without a validation set. Single-line panels annotate their first and last values, and if early_stopping restored the best weights, a dashed line marks that epoch and its value.

Parameters:

Name	Type	Description	Default
history		the History returned by model.fit, or its .history dict.	required
early_stopping		optional EarlyStopping callback; if it restored best weights, that epoch is marked.	None
figsize		optional (width, height); chosen from the number of panels if omitted.	None
ncols		number of columns in the panel grid.	3

Source code in sparsepixels/utils.py

def plot_history(history, early_stopping=None, figsize=None, ncols=3, save_path=None):
    """Plot the training history recorded by SparseTrainingMonitor in one figure.

    Shows the total loss and each compiled metric (train, and validation when present), the loss
    breakdown (task, EBOPS and budget penalties, in red) and the sparse diagnostics (n_max_pixels,
    threshold, ebops, in green). Panels are only created for keys that are present, so it works with
    or without a validation set. Single-line panels annotate their first and last values, and if
    early_stopping restored the best weights, a dashed line marks that epoch and its value.

    Args:
        history: the History returned by model.fit, or its .history dict.
        early_stopping: optional EarlyStopping callback; if it restored best weights, that epoch is marked.
        figsize: optional (width, height); chosen from the number of panels if omitted.
        ncols: number of columns in the panel grid.
    """
    h = history.history if hasattr(history, "history") else dict(history)

    restored = None
    if early_stopping is not None and getattr(early_stopping, "restore_best_weights", False):
        best = getattr(early_stopping, "best_epoch", None)
        if best is not None and best >= 0:
            restored = best + 1

    def _is_single(k):
        return (
            k in ("ebops", "loss_task", "loss_ebops", "loss_n", "loss_maskedE")
            or k.startswith("n_max_pixels")
            or k.startswith("threshold")
        )

    single_keys = [k for k in h if _is_single(k)]
    skip = set(single_keys) | {"lr", "learning_rate"}
    metric_keys = [k for k in h if not k.startswith("val_") and k not in skip]
    m_order = ["loss"]
    metric_keys = [k for k in m_order if k in metric_keys] + [k for k in metric_keys if k not in m_order]
    s_order = ["loss_task", "loss_ebops", "loss_n", "loss_maskedE", "n_max_pixels", "threshold", "ebops"]
    single_keys = [k for k in s_order if k in single_keys] + [k for k in single_keys if k not in s_order]

    panels = metric_keys + single_keys
    if not panels:
        raise ValueError("history has no recorded quantities to plot")

    n = len(panels)
    ncols = min(ncols, n)
    nrows = (n + ncols - 1) // ncols
    fig, axes = plt.subplots(nrows, ncols, figsize=figsize or (5 * ncols, 3.5 * nrows), squeeze=False)
    axes = axes.ravel()

    for ax, key in zip(axes, panels):
        ep = range(1, len(h[key]) + 1)
        show_legend = False
        if key in metric_keys:
            ax.plot(ep, h[key], label=f"train {key}")
            if f"val_{key}" in h:
                ax.plot(ep, h[f"val_{key}"], label=f"val {key}")
            ax.grid(alpha=0.25)
            show_legend = True
        else:
            vals = list(h[key])
            color = "tab:red" if key in ("loss_task", "loss_ebops", "loss_n", "loss_maskedE") else "tab:green"
            ax.plot(ep, vals, marker=".", color=color)
            # annotate first and last epoch (above the point) so the endpoints are easy to read
            for xi, yi in ((1, vals[0]), (len(vals), vals[-1])):
                ax.annotate(
                    f"{yi:g}", (xi, yi), textcoords="offset points", xytext=(0, 6), ha="center", fontsize=8, color=color
                )
            # annotate the restored-epoch value below the point
            if restored is not None and 1 <= restored <= len(vals):
                ax.annotate(
                    f"{vals[restored - 1]:g}",
                    (restored, vals[restored - 1]),
                    textcoords="offset points",
                    xytext=(0, -13),
                    ha="center",
                    va="top",
                    fontsize=8,
                    color="0.35",
                )
            if key == "ebops" and min(vals) > 0:
                ax.set_yscale("log")
            ax.margins(y=0.15)
            ax.grid(True, alpha=0.4)
        if restored is not None:
            ax.axvline(restored, color="0.35", ls="--", lw=1.2, label=f"restored @ epoch {restored}")
            show_legend = True
        if show_legend:
            ax.legend(loc="best", fontsize=8)
        ax.set_xlabel("epoch")
        ax.set_ylabel(key)
        ax.set_title(key)

    for ax in axes[n:]:
        ax.axis("off")
    fig.tight_layout()
    if save_path is not None:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.show()

plot_layer_outputs

plot_layer_outputs(model, x, index=0, layers=None, cmap='viridis', aspect='auto', max_channels=16, save_path=None)

Plot every layer's output for one input image, channel by channel.

Feeds one image through the model and draws one figure per layer, with that layer's channels stacked as rows on a common intensity scale (empty pixels in white) and the active-pixel count annotated per channel. Use it to spot where the sparse signals get over-compressed, e.g. a pooling layer merging the few active pixels of a track into far fewer, and adjust the architecture accordingly.

Parameters:

Name	Type	Description	Default
model		a functional keras model.	required
x		image array of shape (N, H, W, C) or a single image (H, W, C).	required
index		which image of x to feed.	0
layers		layer names to show; defaults to every layer with a plottable output.	None
cmap		colormap for the active values; exact zeros are drawn white.	'viridis'
aspect		imshow aspect for the feature-map panels ("auto" fills each panel, None keeps the natural pixel ratio, a float scales the pixel height).	'auto'
max_channels		cap on the number of channels drawn per layer.	16
save_path		optional path; each layer's figure is saved with the layer name appended.	None

Source code in sparsepixels/utils.py

def plot_layer_outputs(model, x, index=0, layers=None, cmap="viridis", aspect="auto", max_channels=16, save_path=None):
    """Plot every layer's output for one input image, channel by channel.

    Feeds one image through the model and draws one figure per layer, with that layer's channels
    stacked as rows on a common intensity scale (empty pixels in white) and the active-pixel count
    annotated per channel. Use it to spot where the sparse signals get over-compressed, e.g. a
    pooling layer merging the few active pixels of a track into far fewer, and adjust the
    architecture accordingly.

    Args:
        model: a functional keras model.
        x: image array of shape (N, H, W, C) or a single image (H, W, C).
        index: which image of x to feed.
        layers: layer names to show; defaults to every layer with a plottable output.
        cmap: colormap for the active values; exact zeros are drawn white.
        aspect: imshow aspect for the feature-map panels ("auto" fills each panel, None keeps the
            natural pixel ratio, a float scales the pixel height).
        max_channels: cap on the number of channels drawn per layer.
        save_path: optional path; each layer's figure is saved with the layer name appended.
    """
    x = np.asarray(x, dtype="float32")
    if x.ndim == 3:
        x = x[None]
    xb = x[index : index + 1]

    outputs = {}
    for layer in model.layers:
        if layer.__class__.__name__ == "InputLayer":
            continue
        if layers is not None and layer.name not in layers:
            continue
        try:
            out = layer.output
        except Exception:
            continue
        if isinstance(out, (list, tuple)):
            out = out[0]
        outputs[layer.name] = out
    probe = keras.Model(model.inputs, outputs)
    acts = probe.predict(xb, verbose=0)

    cmap_obj = plt.get_cmap(cmap).copy()
    cmap_obj.set_bad("white")

    for name, act in acts.items():
        a = np.asarray(act)[0]
        if a.ndim == 1:
            a = a[None, :, None]
        elif a.ndim != 3:
            continue
        n_ch = min(a.shape[-1], max_channels)
        vmin, vmax = min(0.0, float(a.min())), float(a.max()) or 1.0

        fig, axes = plt.subplots(n_ch, 1, figsize=(14, 0.75 * n_ch + 0.6), squeeze=False)
        for ch in range(n_ch):
            img, ax = a[..., ch], axes[ch][0]
            ax.imshow(
                np.ma.masked_equal(img, 0.0),
                cmap=cmap_obj,
                vmin=vmin,
                vmax=vmax,
                aspect=aspect if aspect is not None else "equal",
                interpolation="nearest",
                origin="lower",
            )
            ax.set_ylabel(f"ch {ch}\n{int((img != 0).sum())} act", fontsize=7, rotation=0, ha="right", va="center")
            ax.set_xticks([])
            ax.set_yticks([])
        shown = f", showing {n_ch} of {a.shape[-1]} channels" if n_ch < a.shape[-1] else ""
        axes[0][0].set_title(f"{name}: output {tuple(np.asarray(act).shape[1:])}{shown}", loc="left", fontsize=10)
        fig.tight_layout()
        if save_path is not None:
            root, ext = os.path.splitext(save_path)
            fig.savefig(f"{root}_{name}{ext or '.png'}", dpi=150, bbox_inches="tight")
        plt.show()

freeze_quantization

freeze_quantization(model, freeze_budget=True)

Freeze the quantization (and optionally the pixel budget) at the current learned values.

Sets every quantizer sublayer non-trainable and zeroes the WRAP integer-part decay, so the bit-width configuration stays exactly fixed afterwards; optionally also freezes the InputReduce budget. Use for two-stage training: after the quantization-aware stage converges, freeze, recompile with a lower learning rate, and fine-tune. The weights then adapt to the final bit assignment instead of a moving one, and early stopping effectively monitors the pure task loss since the frozen regularizer terms are constants.

Parameters:

Name	Type	Description	Default
model		a keras model built with the quantized (HGQ) layers.	required
freeze_budget		also set InputReduce layers non-trainable (a learnable budget n stops moving).	True

Returns:

Name	Type	Description
int		the number of layers set non-trainable. Remember to call model.compile again.

Source code in sparsepixels/utils.py

def freeze_quantization(model, freeze_budget=True):
    """Freeze the quantization (and optionally the pixel budget) at the current learned values.

    Sets every quantizer sublayer non-trainable and zeroes the WRAP integer-part decay, so the
    bit-width configuration stays exactly fixed afterwards; optionally also freezes the InputReduce
    budget. Use for two-stage training: after the quantization-aware stage converges, freeze,
    recompile with a lower learning rate, and fine-tune. The weights then adapt to the final bit
    assignment instead of a moving one, and early stopping effectively monitors the pure task loss
    since the frozen regularizer terms are constants.

    Args:
        model: a keras model built with the quantized (HGQ) layers.
        freeze_budget: also set InputReduce layers non-trainable (a learnable budget n stops moving).

    Returns:
        int: the number of layers set non-trainable. Remember to call model.compile again.
    """
    n_frozen = 0
    for sub in model._flatten_layers(include_self=False):
        cls = sub.__class__
        is_quantizer = "quantizer" in cls.__module__.lower() or "Quantizer" in cls.__name__
        if is_quantizer or (freeze_budget and isinstance(sub, InputReduce)):
            sub.trainable = False
            n_frozen += 1
    for w in model.weights:
        if "i_decay" in w.path:
            w.assign(np.zeros(w.shape))
    return n_frozen

set_sparse_ebops_factor

set_sparse_ebops_factor(model)

Correct each sparse conv's EBOPS for the fact that it only computes on the active pixels.

HGQ counts a conv's EBOPS as if it ran over the whole HW feature map, which overestimates a sparse conv that only touches about n_max_pixels of them. This scales each sparse conv's EBOPS by n_max_pixels / (HW) through HGQ's own ebops_factor, so both the reported EBOPS and the EBOPS regularization loss reflect the sparse compute. SparseTrainingMonitor applies this automatically at the start of training, so you usually do not need to call it directly -- call it manually only to correct EBOPS you inspect before training, or to refresh it if a learnable budget moves a lot (it uses the current n_max_pixels, a good approximation since the budget moves slowly).

Returns a dict of {layer_name: factor} for the sparse convs that were adjusted.

Source code in sparsepixels/utils.py

def set_sparse_ebops_factor(model):
    """Correct each sparse conv's EBOPS for the fact that it only computes on the active pixels.

    HGQ counts a conv's EBOPS as if it ran over the whole H*W feature map, which overestimates a
    sparse conv that only touches about n_max_pixels of them. This scales each sparse conv's EBOPS by
    n_max_pixels / (H*W) through HGQ's own ebops_factor, so both the reported EBOPS and the EBOPS
    regularization loss reflect the sparse compute. SparseTrainingMonitor applies this automatically
    at the start of training, so you usually do not need to call it directly -- call it manually only
    to correct EBOPS you inspect before training, or to refresh it if a learnable budget moves a lot
    (it uses the current n_max_pixels, a good approximation since the budget moves slowly).

    Returns a dict of {layer_name: factor} for the sparse convs that were adjusted.
    """
    scale = _sparse_conv_scale(model)
    applied = {}
    for layer in model.layers:
        if isinstance(layer, QConv2DSparse) and id(layer.conv) in scale:
            layer.conv.ebops_factor = scale[id(layer.conv)]
            applied[layer.name] = scale[id(layer.conv)]
    return applied

Quantization diagnostics

print_quantization

print_quantization(model)

Print the per-layer quantizer bit-widths (kernel, bias, input as mean, min, max) and EBOPS.

The EBOPS column reads HGQ's per-layer value, which reflects each layer's ebops_factor, so it is sparse corrected once set_sparse_ebops_factor has been applied.

Source code in sparsepixels/utils.py

def print_quantization(model):
    """Print the per-layer quantizer bit-widths (kernel, bias, input as mean, min, max) and EBOPS.

    The EBOPS column reads HGQ's per-layer value, which reflects each layer's ebops_factor, so it is
    sparse corrected once set_sparse_ebops_factor has been applied.
    """
    print(f"\nModel: {model.name}")
    header = (
        f"{'Layer':<12} {'#Kernel':>8} {'#Bias':>6}"
        f"  {'K mean':>6} {'min':>5} {'max':>5}"
        f"  {'B mean':>6} {'min':>5} {'max':>5}"
        f"  {'I mean':>6} {'min':>5} {'max':>5}"
        f"  {'eBOPs':>10}"
    )
    print(header)
    print("-" * len(header))
    total_ebops = 0.0
    for layer in model.layers:
        info = _get_layer_info(layer)
        if info is None:
            continue
        km, klo, khi = _fmt_bits(info["kq"].bits) if info["kq"] else ("-", "-", "-")
        bm, blo, bhi = _fmt_bits(info["bq"].bits) if info["bq"] else ("-", "-", "-")
        im, ilo, ihi = _fmt_bits(info["iq"].bits) if info["iq"] else ("-", "-", "-")
        eb = info["ebops"]
        if eb is not None:
            total_ebops += eb
            es = f"{eb:.0f}"
        else:
            es = "-"
        print(
            f"{info['name']:<12} {info['n_kernel']:>8} {info['n_bias']:>6}"
            f"  {km:>6} {klo:>5} {khi:>5}"
            f"  {bm:>6} {blo:>5} {bhi:>5}"
            f"  {im:>6} {ilo:>5} {ihi:>5}"
            f"  {es:>10}"
        )
    print("-" * len(header))
    print(f"{'Total eBOPs':>{len(header) - 10}}{total_ebops:>10.0f}")

plot_quantization

plot_quantization(models, figsize=(14, 3), save_path=None)

Violin plots of the per-layer kernel, bias and input bit-width distributions.

Parameters:

Name	Type	Description	Default
models		a model or list of models to compare (one row per model).	required
figsize		(width, per-model height) of the figure.	(14, 3)

Source code in sparsepixels/utils.py

def plot_quantization(models, figsize=(14, 3), save_path=None):
    """Violin plots of the per-layer kernel, bias and input bit-width distributions.

    Args:
        models: a model or list of models to compare (one row per model).
        figsize: (width, per-model height) of the figure.
    """
    if not isinstance(models, (list, tuple)):
        models = [models]
    categories = [("kq", "Kernel bits"), ("bq", "Bias bits"), ("iq", "Input bits")]
    fig, axes = plt.subplots(
        len(models),
        len(categories),
        figsize=(figsize[0], figsize[1] * len(models)),
        squeeze=False,
        constrained_layout=True,
    )
    for row, model in enumerate(models):
        infos = [info for layer in model.layers if (info := _get_layer_info(layer)) is not None]
        colors = plt.cm.tab10(np.linspace(0, 1, max(len(infos), 1)))
        for col, (key, title) in enumerate(categories):
            ax = axes[row][col]
            names, data, used_colors = [], [], []
            for info, color in zip(infos, colors):
                if info[key] is None:
                    continue
                names.append(info["name"])
                data.append(np.array(info[key].bits).flatten())
                used_colors.append(color)
            if not data:
                ax.text(0.5, 0.5, "N/A", transform=ax.transAxes, ha="center", va="center", fontsize=14, color="gray")
                ax.set_title(f"{model.name} - {title}")
                continue
            parts = ax.violinplot(data, positions=range(len(data)), vert=False, showmedians=True, showextrema=False)
            for body, c in zip(parts["bodies"], used_colors):
                body.set_facecolor(c)
                body.set_alpha(0.7)
            parts["cmedians"].set_color("black")
            ax.set_yticks(range(len(names)))
            ax.set_yticklabels(names)
            ax.set_xlabel("Bitwidth")
            allb = np.concatenate([d.ravel() for d in data])
            lo, hi = float(allb.min()), float(allb.max())
            pad = max(0.5, 0.08 * (hi - lo))  # show the full band with a small margin at both ends
            ax.set_xlim(lo - pad, hi + pad)
            ax.set_title(f"{model.name} - {title}")
    if save_path is not None:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.show()