Hardening the MobileNet TFLite Toolkit (Medium-safe test)

This is a follow-up to my two earlier articles, Creating Quantized Multi-Task MobileNet V3 Models for Edge Deployment and From Multi-Task to Universal. Together they introduced a tool that takes a single CLI invocation and emits a quantized multi-head MobileNet ready for an edge device. This article covers the changes that landed since — the unglamorous but important work of taking a demo-shaped tool and making it something you can actually depend on in a deployment.

TL;DR. Two merged branches added five things to the toolkit: a --backbone flag covering MobileNet V1/V2/V3/V4 plus a custom V3-Small ( v3c) for grayscale and odd input shapes; per-head feature taps with the --heads "N@stride" syntax; opt-in FPN fusion with --fusion fpn; real-image calibration via --calibration-data-dir; and a 78-test pytest suite that, together with 20 end-to-end example scripts, runs as a 98-check gate before every change.

The gap between a demo and a deployment

The earlier articles ended with a tool that could generate one kind of model — quantized, multi-head, MobileNetV3-based — and 12 head types covering most everyday CV tasks. That was enough to write a blog post about. It wasn’t enough to deploy a fleet on.

Five things stood out as “demo-shaped” rather than “production-shaped”:

• Only one backbone. V3 is great, but it’s not the only MobileNet, and it can’t always be the right pick. Some chips have V2 or V1 kernels deeply optimized; some tasks want V4’s MQA blocks; some pipelines want a custom V3-Small that handles grayscale.
• Quantization calibrated on random noise. Fine for verifying shapes and op coverage; bad for the actual int8 scales the model will run with on real inputs.
• Every head saw the same feature map. Classification doesn’t need fine spatial resolution, but detection and segmentation do. There was no way to point different heads at different backbone strides.
• No cross-scale fusion. Even with per-head taps, each head saw only its own stride. There was no path to combine semantic depth with spatial precision.
• No test suite. The CLI worked when I ran it. That’s not a guarantee — that’s a hope.

Two branches ( all-mnvx then advanced-features) merged sixteen commits' worth of work into main to close those gaps. The rest of this article walks through what changed and why each choice came out the way it did.

Five backbones, one CLI: --backbone v1 | v2 | v3 | v4 | v3c

The first change was the smallest in lines and the largest in scope. The CLI used to be called create_quantized_mobilenet_v3.py. Now it's create_quantized_mobilenet.py with a --backbone flag that dispatches to one of five architecture classes. Everything else - heads, quantization, calibration, naming - is unchanged across backbones.

python src/create_quantized_mobilenet.py \ 
    --backbone v1 \ 
    --alpha 0.25 \ 
    --input-shape "96x96x1" \ 
    --heads "5,2,3" \ 
    --head-names "object_class,person_detection,age_group" \ 
    --output-name mnv1_multi

The architecture-class dispatch lives in the CLI helper that resolves a backbone name to its class. Each backbone has its own multi-head QAT architecture ( MultiHeadMobileNetV1QATArchitecture, V2, V4, plus the original V3) so the head-attachment logic is shared while backbone-specific quirks stay encapsulated.

The fifth option, v3c, is the most interesting. It wires in a project-local V3-Small implementation rather than reaching for tf.keras.applications.MobileNetV3Small. Why? Because Keras-applications V3 hard-codes some input assumptions - three channels, certain shapes - that fall apart on real edge work. Grayscale sensors and 64×48 thermal frames are the rule on small devices, not the exception.

The five backbones, and when to reach for each:

• **v1** - origin: Keras-applications. Smallest kernels, deeply optimized everywhere. Best for MCU-class devices with a strict ops budget.
• **v2** - origin: Keras-applications. Balanced; inverted residuals; mature tooling. A safe default for older accelerators.
• **v3** - origin: Keras-applications. Hard-swish + SE blocks; better accuracy per unit cost. The default for newer NPUs.
• **v4** - origin: Keras-applications. MQA blocks; the most recent design. For the latest accelerators that support it.
• **v3c** - origin: project-local custom. Grayscale and odd shapes; full control. For thermal, depth, and single-channel sensors.

The “right” backbone is rarely a model-quality question — it’s a question of which kernels the target chip ships with. Having all five in one CLI means picking one is a --backbone flag away.

Real calibration: --calibration-data-dir IMG_DIR

Post-training int8 quantization needs a representative dataset. The TFLite converter feeds that dataset through the float model to observe the dynamic range of every activation tensor, then picks scales and zero-points that map those ranges onto int8.

The demo version of the tool fed it random uniform noise. That’s enough to verify the converter runs and the output graph is shape-correct. It is not enough to deploy.

Random noise has a flat, dense histogram across [0, 1). Real images don't. They have black backgrounds, sky regions, skin tones, edges, textures - distributions that bunch up around specific values and ignore others entirely. When the calibrator sees only noise, it sizes the int8 range to a distribution your real inputs never produce. Quantization error gets concentrated exactly where your data actually lives.

The fix is a flag:

python src/create_quantized_mobilenet.py \ 
    --backbone v3 \ 
    --alpha 0.25 \ 
    --input-shape "96x96x3" \ 
    --heads "5,2,3" \ 
    --calibration-samples 50 \ 
    --calibration-data-dir /path/to/held_out_validation_images

The directory is scanned for PNG/JPG files; up to --calibration-samples are loaded, decoded, resized, and fed to the converter. If the flag is omitted, the behavior is unchanged - random noise still works for shape verification or quick smoke tests.

Random calibration:

• Setup: none.
• What it verifies: converter runs, shapes correct.
• When to use: development, shape checks, CI.
• Accuracy impact: unpredictable, often poor.

Real-image calibration:

• Setup: a directory of representative images.
• What it verifies: converter runs and int8 scales fit the data.
• When to use: every deployable model.
• Accuracy impact: matches the input distribution.

The point isn’t that random calibration is broken. It’s that the calibration step is the only thing standing between a float model and an int8 model, and using realistic data for it is free.

Per-head feature taps: --heads "5@32,2@16,3@8"

The original tool assumed every head was a classification head, so it stacked them all on the final feature map at the backbone’s deepest stride (stride 32 for a 96×96 input). That assumption breaks for almost everything else.

A 1×1×N classification head wants depth - semantic content. A segmentation head wants spatial resolution. A detection head wants both, and ideally at multiple scales. Forcing every head onto the deepest map throws away the early- and mid-stage feature maps that the backbone already computed.

The new --heads syntax accepts a stride per head:

--heads "5@32,2@16,3@8" 
--head-names "object_class,person_detection,age_group"

The classification head taps stride 32 (semantically rich, low spatial resolution). The detection head taps stride 16 (mid-scale). The age head taps stride 8 (fine spatial). Each backbone exposes its own _features_by_stride mapping so the CLI can route correctly across V1, V2, V3, V4, and V3c.

The bare form (--heads "5,2,3") still works and behaves exactly as before - every head taps the default deepest stride. The @N form is opt-in, so old example scripts and old config files keep working.

A companion change, add_head_dynamically(tap_stride=...), makes it possible to attach a new head to an already-trained, frozen backbone - useful when a downstream team needs to add a task without retraining the shared trunk.

Cross-scale fusion: --fusion fpn

Per-head taps fix the “wrong scale” problem but introduce a new one: each head still sees its stride in isolation. A head at stride 8 has fine spatial precision but shallow semantic content; a head at stride 32 has rich semantics but no spatial precision. Tasks that want both — small-object detection, fine segmentation — get neither.

The classic answer is a Feature Pyramid Network. The new --fusion fpn flag builds one on top of the existing feature-tap mechanism. For each backbone stride that any head uses, a 1×1 lateral convolution projects it to a shared channel count (--fpn-channels, default 64). Then a top-down pathway upsamples each level by 2× and adds it to the next-finer lateral. The result is a pyramid where every level is both semantically deep and spatially precise.

python src/create_quantized_mobilenet.py \ 
    --backbone v3 \ 
    --alpha 0.25 \ 
    --input-shape "96x96x3" \ 
    --heads "5@32,2@16,3@8"\ 
    --fusion fpn \ 
    --fpn-channels 64

FPN is opt-in for a reason. It adds parameters (~ fpn_channels × levels worth of lateral convs plus the top-down adds) and adds depth to the int8 graph, which the calibrator and the on-device runtime both have to swallow. For pure classification, it's overkill; the default is still single-tap (--fusion none). For detection and segmentation, the parameter cost usually pays for itself in accuracy.

The 98-check gate

The single most valuable change in the advanced-features branch wasn't a feature - it was a test suite. Before the branch, the CLI worked when I ran it. After, every change runs through 98 checks: 78 pytest unit and integration tests, plus 20 end-to-end example scripts driven by test_all_examples.sh.

tests/ 
├── test_architectures.py (126 lines) 
├── test_calibration.py (93 lines) 
├── test_cli_parsing.py (76 lines) 
├── test_dynamic_head.py (63 lines) 
├── test_fpn.py (103 lines) 
├── test_head_configuration.py (98 lines) 
├── test_losses.py (143 lines) 
└── test_reporting.py (38 lines)

The split between pytest and example scripts is deliberate. Pytest covers the surface — HeadConfiguration validation, the new N@stride parser, the per-backbone _features_by_stride lookups, calibration loading, dynamic head addition, the rewritten losses, FPN wiring. The example scripts cover what's actually shipped: every CLI flag combination that appears in the README runs end-to-end to a verified TFLite file.

The gate has two layers:

• **pytest tests/ - catches API contracts, parser edge cases, single-unit correctness. 78 checks.**
• **bash test_all_examples.sh - catches CLI integration, real TFLite emission, shape verification. 20 checks.**
• Total gate — everything between a commit and a deploy. 98 checks.

Two loss functions got rewrites as part of the same branch. YOLOLoss is now num_classes-aware with proper objectness masking - the previous implementation conflated classes and confidence. TextDetectionLoss accepts both a dict (when used programmatically) and a single concatenated tensor (the shape Keras's model.compile() actually passes at runtime). Both rewrites came with their own pytest cases in test_losses.py.

Trade-offs

None of these changes are free.

• Five backbones means five sets of edge cases. The architecture-class dispatch is small, but _features_by_stride had to be implemented per backbone - V1, V2, V3, V4, and V3c each name their layers differently. The custom v3c adds 267 lines of architecture code that no upstream library maintains for me.
• Real calibration shifts the trust model. Once a directory of images is part of the build, it has to be versioned, secured, and kept representative. A stale calibration set quietly produces a worse model every time the data distribution drifts.
• FPN adds parameters and depth. For an fpn_channels=64 pyramid over three levels, you add roughly 3 lateral 1×1 convolutions plus the top-down adds. On an MCU-class target, that overhead matters.
• **v3c only handles MobileNetV3-Small variants.** The custom backbone covers the grayscale and odd-shape gap, but it does not generalize to V3-Large or V1/V2. Adding more custom backbones would cost real engineering time.
• A 98-check gate is a 98-check wait. Running the full suite locally takes minutes, not seconds. Pre-commit becomes a thing you have to think about.

These are the costs of the toolkit being deployable rather than demonstrable. They feel worth paying.

Takeaways

• The hard work in shipping a model toolkit isn’t usually in the model. It’s in the seams around the model: which architecture, which calibration data, which features go to which head, and what evidence you have that any of it still works.
• A --backbone flag is a tiny change in the CLI surface and a large change in what hardware you can target.
• Random calibration is a useful debugging affordance and a terrible default. The fix is one flag.
• Feature taps and FPN are independent tools. Taps say where each head reads. FPN says with what context. Most projects need one before they need the other.
• The test suite isn’t a deliverable — it’s the only mechanism that lets the other deliverables stay correct as the codebase grows.

The code is open source at github.com/saeid-h/universal-multitask-mobilenet-tflite. Each feature in this article corresponds to one example script in examples/ - 13 through 20 cover the new flags end-to-end.