Initial release: technical documentation, quantization scripts, and NPU benchmarks

73072fb verified 3 days ago

17.1 kB

METHODOLOGY: W4A8 Quantization Pipeline

This document details the technical steps, tools, and configurations used to generate the 12 quantized models optimized for the Snapdragon X Elite NPU.

1. Architectural Overview

The target models are:

MobileNet V4-S (Scale): Lightweight model optimized for high-speed metrology.
MobileNet V4-L (Rocks): Large model for higher geological precision.
MobileNet V5-300M (Rocks): Latest MobileNet architecture utilizing GELU activations and optimized bottleneck layers (analogous to the Gemma 3 Nano vision encoder).

2. Quantization Formats & Zero-FP32 Policy

Two specific formats were produced for each architecture to support the tiered deployment strategy, governed by a strict Zero-FP32 Hardware Constraint:

INT8 Full Integer (A8W8): 8-bit Activations, 8-bit Weights. Internal operations are pure integer kernels (QLinearConv, QLinearMul, QLinearAdd). Inputs and Outputs forced to INT8.
W4A8 Full Integer (A8W4): 8-bit Activations, 4-bit Weights (packed sub-byte integers). Inputs and Outputs forced to INT8 with 0 dynamic float32 dequantization layers.

Zero-FP32 Qualcomm Hexagon HTP Requirement: Qualcomm Hexagon HTP accelerators immediately reject or fracture graphs containing mixed floating-point tensors in the convolutional backbone. We enforce:

ONNX: Exporté via onnxruntime.quantization.quantize_static avec QuantFormat.QDQ (QuantType.QInt8 / QuantType.QUInt4), combiné à la transformation post-traitement force_full_integer_io. Cette dernière élimine les nœuds d'entrée/sortie QuantizeLinear et DequantizeLinear pour exposer un graphe dont toutes les entrées/sorties et tenseurs intermédiaires sont strictement des entiers (int8_t). Ce format QDQ élagué est directement compilé en nœuds matériels natifs sans casting flottant.

TFLite: Full Integer compilation where supported_ops = [TFLITE_BUILTINS_INT8] and all depthwise/convolutional layers are bound to integer kernels with 0 FP32 buffers (0.0% float tensors).

3. Mathematical Foundations

The quantization uses linear symmetric/asymmetric mapping. The relationship between a floating-point value $q_f$ and its quantized integer value $q_i$ is:

$q_f = \text{scale} \times (q_i - \text{zero\_point})$

Where:

Scale: Floating-point scaling factor.
Zero Point: Integer offset (typically 0 or -128 for INT8).

For W4A8 models:

Weights (W4): Weights are grouped and quantized block-wise to minimize precision loss in deep convolutional layers.
Activations (A8): Activations use static calibration (observed min-max) to minimize saturation of outlier values.

MobileNet V4 (S & L) Pipeline:

Tool: ai-edge-quantizer (Google).
Process:
- Conversion from SavedModel to FP32 TFLite with Signatures.
- Static calibration on 100 representative geological samples.
- Export with Full Integer flag to ensure zero floating-point operators.
I/O Optimization: Used the ai-edge-quantizer API to specify input_type=int8 and output_type=int8.

MobileNet V5 Specifics:

Challenge: Presence of GELU (Erf) operations and memory-intensive tracking of 300M parameters.
NPU Solution: Implemented Surgical Graph Optimization to replace Erf with a polynomial approximation ($x \times \sigma(1.702 \times x)$).
Graph Freezing: To ensure the model is fully compatible with LiteRT without triggering _DictWrapper or untracked resource errors during conversion, all 294M variables were frozen into constants directly in the graph definition.
4-bit Quantization: Executed via the ai-edge-quantizer compression API, allowing weight reduction to 4-bits while maintaining 8-bit activations for NPU compatibility. This results in a "No-Flex" artifact where 100% of the operators are natively supported by the NPU.

4. ONNX Quantization Details (Static QDQ & Stripped Full Integer I/O)

General Pipeline:

FP32 ONNX -> ONNX Runtime Static Quantization (QuantFormat.QDQ) -> force_full_integer_io (Stripping FP32 Casts).

The quantization workflow uses onnxruntime.quantization.quantize_static with QuantFormat.QDQ:

QuantFormat.QDQ Pipeline: Activations and weights are quantized using QuantizeLinear and DequantizeLinear pairs configured for QuantType.QInt8 (or QUInt4 for 4-bit weights).
Post-Processing (Full Integer I/O Stripping): Standard QDQ models expose FP32 inputs and outputs. To eliminate CPU-bound float conversion layers on edge devices, the custom routine force_full_integer_io rewrites the model graph by removing boundary quantize/dequantize nodes and updating graph signatures to native int8_t inputs/outputs.
NPU Compatibility: The resulting graph presents pure integer boundaries directly consumable by execution providers (such as QNN Execution Provider) with zero partition breaks.

Calibration Data:

Scale: calib_echelle.npy (100 samples, NCHW [1, 3, 1024, 1024]).
Rocks V4: calib_roches_244.npy (100 samples, NCHW [1, 3, 244, 244]).
Rocks V5: calib_roches_256.npy (100 samples, NCHW [1, 3, 256, 256], transposed to NHWC).

Forcing Full Integer I/O:

Standard ONNX tools insert QuantizeLinear (input) and DequantizeLinear (output) nodes that expect floats. To eliminate CPU-bound type casting:

A custom post-processing script stripped these boundary nodes.
Graph metadata was manually adjusted to declare INT8 as the native I/O type.

Calibration Dataset Preparation

To ensure bit-accurate min-max profiling, we prepared a calibration dataset of 100 representative samples. The following script was used to generate the .npy tensors:

import os, glob
import numpy as np
from PIL import Image

def prepare_calib(image_dir, output_path, size=(256, 256), num_samples=100):
    all_files = []
    for ext in ['*.jpg', '*.png', '*.jpeg']:
        all_files.extend(glob.glob(os.path.join(image_dir, "**", ext), recursive=True))
    
    np.random.seed(42)
    np.random.shuffle(all_files)
    selected = all_files[:num_samples]
    
    images = []
    for f in selected:
        img = Image.open(f).convert('RGB').resize(size, Image.Resampling.LANCZOS)
        img_data = np.array(img).astype(np.float32) / 255.0
        # NCHW format
        img_data = np.transpose(img_data, (2, 0, 1)) 
        images.append(img_data)
            
    np.save(output_path, np.stack(images))

# Usage for MobileNetV5 (256x256)
# prepare_calib("path/to/images", "calib_roches_256.npy", size=(256, 256))

5. Calibration & Preprocessing

Data Selection:

We selected 100 images per model from the validation datasets.

Representativeness: Images cover various lighting conditions and rock textures to ensure the min/max calculation during calibration is robust (minimizing KL Divergence).
Source Format: .npy (Numpy) files containing normalized tensors.

Parameter	Scale V4-S	Rocks V4-L	Rocks V5
Resolution	1024 x 1024	244 x 244	256 x 256
ONNX Format	NCHW (1,3,1024,1024)	NCHW (1,3,244,244)	NHWC (1,256,256,3)
Normalization	[0, 1]	[0, 1]	[0, 1]

6. Edge Hardware Optimization & Portability

While benchmarks were conducted on the Snapdragon X Elite, the optimization strategy was designed for maximum cross-platform portability:

W4A8 Format Selection: The W4A8 format was chosen over more specialized formats (like W4A16) to ensure the widest possible compatibility across various Edge AI runtimes (LiteRT, ONNX Runtime) and hardware accelerators (NPUs from multiple vendors, modern GPUs). This ensures a consistent performance baseline across the Windows ARM64 Copilot+ ecosystem.
Full Integer I/O: Without this, the CPU would perform a "Casting" overhead for every frame. By forcing INT8, the data flow remains entirely in the integer domain from the sensor to the final prediction, significantly reducing thermal throttling risks on mobile devices.

7. Troubleshooting

A. Erf / GELU (MobileNet V5)

Erf is not natively supported by most 8-bit accelerators.

Solution: Removed SELECT_TF_OPS in favor of a native subgraph approximation. This ensures that the Qualcomm Hexagon NPU can execute the entire model without falling back to the CPU for activation kernels.

B. Mismatch in Dimensions

V4 and V5 models use different conventions (NCHW vs NHWC).

Solution: Implemented a dynamic transpose=(0, 2, 3, 1) in the CalibrationDataReader to adapt on-the-fly without modifying source .npy files.

8. Technical Implementation (Python API)

A. TFLite Quantization (LiteRT)

Using ai-edge-quantizer:

1. Rock Classification (MobileNetV5 - Gemma 3 Nano)

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

def export_to_noflex_tflite(keras_model, output_path):
    # 1. Capture model in a concrete function
    @tf.function(input_signature=[tf.TensorSpec([1, 256, 256, 3], tf.float32, name="input_1")])
    def serve(x):
        return {"output_0": keras_model(x, training=False)}

    # 2. Freeze 300M variables into constants (No-Flex requirement)
    frozen_func = convert_variables_to_constants_v2(serve.get_concrete_function())
    
    # 3. Convert via SavedModel with explicit signatures
    temp_sm = "temp_sm"
    tf.saved_model.save(tf.Module(), temp_sm, signatures={"serving_default": frozen_func})
    
    converter = tf.lite.TFLiteConverter.from_saved_model(temp_sm)
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
    tflite_model = converter.convert()
    
    with open(output_path, "wb") as f:
        f.write(tflite_model)

2. W4A8 Compression (Weights 4-bit)

Using the ai_edge_quantizer API for high-density compression:

compression_config = ai_edge_quantizer.CompressionConfig(
    weight_bits=4,
    activation_bits=8,
    asymmetric=True
)
quantizer.export(
    output_path,
    compression_config=compression_config,
    input_type=ai_edge_quantizer.DType.INT8,
    output_type=ai_edge_quantizer.DType.INT8
)

B. Custom Graph Stripping (ONNX Full Integer)

Python logic used to prune dequantization artifacts and force INT8 I/O:

import onnx

def force_full_integer_io(model_path):
    m = onnx.load(model_path)
    g = m.graph
    # Input Node Processing
    inp = g.input[0]
    for n in list(g.node):
        if n.op_type == 'QuantizeLinear' and inp.name in n.input:
            out_name = n.output[0]
            g.node.remove(n)
            for c in g.node:
                for i, in_c in enumerate(c.input):
                    if in_c == out_name: c.input[i] = inp.name
            inp.type.tensor_type.elem_type = 3 # INT8
            break
    # Output Node Processing
    out = g.output[0]
    for n in list(g.node):
        if n.op_type == 'DequantizeLinear' and out.name in n.output:
            in_name = n.input[0]
            g.node.remove(n)
            for p in g.node:
                for i, out_p in enumerate(p.output):
                    if out_p == in_name: p.output[i] = out.name
            out.type.tensor_type.elem_type = 3 # INT8
            break
    onnx.save(m, model_path)

9. Final I/O Specifications

To ensure bit-accurate inference on edge accelerators, all models adhere to the following I/O contracts.

Model	Format	Input Shape	Output Shape	Axis Order
Scale V4-S	TFLite	`[1, 1024, 1024, 3]`	`[1, 1]`	NHWC
Scale V4-S	ONNX	`[1, 3, 1024, 1024]`	`[1, 1]`	NCHW
Rocks V4-L	TFLite	`[1, 256, 256, 3]`	`[1, 104]`	NHWC
Rocks V4-L	ONNX	`[1, 3, 256, 256]`	`[1, 104]`	NCHW
Rocks V5	TFLite	`[1, 256, 256, 3]`	`[1, 104]`	NHWC
Rocks V5	ONNX	`[1, 256, 256, 3]`	`[1, 104]`	NHWC

Input Normalization

All inputs are expected as Signed INT8 tensors in the range [-128, 127].

# Formula for bit-accurate preprocessing (NumPy)
# Maps [0, 1] range to [-128, 127] signed INT8
input_int8 = (pixel_float * 255 - 128).astype(np.int8)

10. Advanced Inference: 21-Pass Multi-Scale Tiling (21-MST)

To prevent the loss of micro-mineralogical information (like small crystals) during downscaling to 256x256, the Lithotheque engine implements a 21-Pass Multi-Scale Tiling strategy.

Instead of a single prediction, the model evaluates 21 distinct crops of the high-resolution input:

1 Global Pass: Full resize (20% weight).
4 Medium Passes: 2x2 grid crops (30% total weight).
16 Fine Passes: 4x4 grid crops (50% total weight).

Reference Python Implementation (scripts/litho_reference_engine.py)

11. Confidence Calibration: Super-Selectivity Matrix

To ensure consistent behavior across diverse hardware tiers and quantization formats, the inference engine applies a dynamic Super-Selectivity Matrix during post-processing. This matrix adjusts the Softmax temperature (multiplier) based on the model's outputScale, compensating for precision loss in highly compressed (Legacy/Balanced) or ultra-precise (Premium) artifacts.

Tier / Model Profile	Scale Threshold	Multiplier	Rationale
Premium / App Internal	`< 0.0003`	1.500f	High precision, slight sharpening of peaks.
Standard Rocks	`< 0.0050`	1.200f	Balanced calibration for INT8.
Probe INT8	`< 0.0100`	1.000f	Identity (raw model selectivity).
Balanced Rocks	`< 0.0240`	0.800f	Compensates for INT4 quantization noise.
Legacy Rocks	`< 0.0300`	0.700f	Maximum smoothing for entry-level devices.
Premium Scale	`< 0.0540`	5.000f	High selectivity for geometric metrology.
Std/Legacy Scale	`< 0.0600`	4.500f	Metric precision for fallback tiers.
Default / Unknown	`else`	4.000f	Safety net for new/unidentified models.

12. Technical Reproduction Scripts

Methodological Blueprint & Adaptability The scripts provided in the scripts/ directory serve as a technical blueprint. The specific resolutions (e.g., 256x256, 1024x1024), axis orderings, and normalization logic used are arbitrary examples based on the Lithotheque engine. Researchers are expected to adapt these scripts (specifically input dimensions and normalization layers) to match their own specific models and datasets.

For full scientific reproducibility, the following specialized automation scripts are provided in the scripts/ directory:

TFLite / LiteRT (ai-edge-quantizer)

MobileNetV5 Rocks: tflite_v5_rocks_int8.py / tflite_v5_rocks_int4.py
MobileNetV4-L Rocks: tflite_v4l_rocks_int8.py / tflite_v4l_rocks_int4.py
MobileNetV4-S Scale: tflite_v4s_scale_int8.py / tflite_v4s_scale_int4.py

ONNX (onnxruntime-quantization)

MobileNetV5 Rocks: onnx_v5_rocks_int8.py / onnx_v5_rocks_int4.py
MobileNetV4-L Rocks: onnx_v4l_rocks_int8.py / onnx_v4l_rocks_int4.py
MobileNetV4-S Scale: onnx_v4s_scale_int8.py / onnx_v4s_scale_int4.py

Multi-Vendor AOT Compilation & Benchmarking

AOT Compilers: compile_qairt_aot.py, compile_tensor_aot.py, compile_mediatek_aot.py, compile_intel_aot.py
NPU Benchmark Harnesses: benchmark_npu_onnx.py, benchmark_npu_qairt.py, benchmark_snapdragon_native_binaries.py

13. Replication Environment

Software Stack

Python: Python 3.11+ (x64 & ARM64)
Frameworks: tensorflow, ai-edge-quantizer, ai-edge-litert, onnxruntime-quantization, torch, torchao.
Compiler / Runtimes: MSVC, CUDA Toolkit, cuDNN, Qualcomm QAIRT SDK.

Hardware (Generation & Quantization Workstation)

GPU: NVIDIA Blackwell Architecture GPU (Compute Capability sm_120, GDDR7)
CPU: Multi-core x86_64 Host Processor (AVX2)
RAM: High-capacity System Memory (32 GB+)
OS: 64-bit Desktop Environment (Windows 11 / Linux WSL2 for AOT Compilers)

Target Edge Benchmark Hardware (Snapdragon Edge Deployment Platform)

Platform: Qualcomm Snapdragon® X Series Reference Platform (ARM64)
NPU: Qualcomm® Hexagon™ HTP NPU (45 TOPS INT8)
RAM: Unified LPDDR5x High-Bandwidth Memory
OS: Windows 11 ARM64 Environment

This document certifies the technical reproducibility of the models across NVIDIA Blackwell and Qualcomm Snapdragon platforms.