Quantization Strategies: Mastering the Bit-Budget

import numpy as np
import onnxruntime
from onnxruntime.quantization import CalibrationDataReader

class ImageNetDataReader(CalibrationDataReader):
    def __init__(self, calibration_image_folder, input_name):
        self.image_folder = calibration_image_folder
        self.input_name = input_name
        self.data_iter = None
        self.preprocessed_data = []
        self.preprocess_images()

    def preprocess_images(self):
        # Load 100-200 images for a good statistical sample.
        # CRITICAL: This MUST match your C++ preprocessing (Mean/StdDev) exactly!
        for img_path in list_images(self.image_folder)[:100]:
            img = load_image(img_path)
            img = resize_aspect_aware(img, (224, 224))
            # ImageNet Normalization
            img = (img / 255.0 - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
            img = img.transpose(2, 0, 1) # HWC to NCHW
            img = np.expand_dims(img, axis=0).astype(np.float32)
            self.preprocessed_data.append({self.input_name: img})

    def get_next(self):
        if self.data_iter is None:
            self.data_iter = iter(self.preprocessed_data)
        return next(self.data_iter, None)

from onnxruntime.quantization import quantize_static, QuantType, QuantFormat

model_fp32 = "models/mobilenetv2.onnx"
model_int8 = "models/mobilenetv2.quant.onnx"

# Create our data reader
dr = ImageNetDataReader("data/calibration_images", "input")

# Run the quantization
quantize_static(
    model_fp32,
    model_int8,
    dr,
    quant_format=QuantFormat.QDQ, # Quantize-Dequantize format (modern ORT standard)
    per_channel=True,             # High precision for depthwise convs
    weight_type=QuantType.QInt8,
    activation_type=QuantType.QInt8,
    nodes_to_exclude=[]           # Add names of sensitive layers here if needed
)
print(f"Quantized model saved to {model_int8}")

Ort::SessionOptions sessionOptions;

// 1. Enable Graph Optimizations
sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

// 2. Append the WebGPU or DirectML provider
// These providers will detect the INT8 nodes and map them to DP4A/Tensor Core instructions
sessionOptions.AppendExecutionProvider("webgpu");

// 3. Load the quantized model
Ort::Session session(env, L"models/mobilenetv2.quant.onnx", sessionOptions);

# Example: Keep the final classifier head in FP32
quantize_static(..., nodes_to_exclude=['Classifier_Head_FC'])

// Logic of the instruction
int dotProduct4x8bitPackedEXT(uint a, uint b, int acc) {
    int res = acc;
    res += int8_t(a & 0xFF) * int8_t(b & 0xFF);
    res += int8_t((a >> 8) & 0xFF) * int8_t((b >> 8) & 0xFF);
    res += int8_t((a >> 16) & 0xFF) * int8_t((b >> 16) & 0xFF);
    res += int8_t((a >> 24) & 0xFF) * int8_t((b >> 24) & 0xFF);
    return res;
}

struct QuantizedLayer {
    std::vector<int8_t> weights; // INT8 weights
    std::vector<float> rescaleFactors; // Pre-calculated (S_in * S_w) / S_out
    int32_t bias_scale; // Usually quantized to 32-bit
};

QuantizedLayer quantizeWeights(const float* fp32Weights, uint32_t outC, uint32_t inC, uint32_t kSize, float inputScale, float outputScale) {
    QuantizedLayer layer;
    uint32_t filterSize = inC * kSize * kSize;
    layer.weights.resize(outC * filterSize);
    layer.rescaleFactors.resize(outC);

    for (uint32_t oc = 0; oc < outC; ++oc) {
        // 1. Find max for this channel
        float maxW = 0.0f;
        for (uint32_t i = 0; i < filterSize; ++i) {
            maxW = std::max(maxW, std::abs(fp32Weights[oc * filterSize + i]));
        }

        // 2. Calculate Weight Scale
        float weightScale = maxW / 127.0f;

        // 3. Pre-calculate the Rescale Factor for the shader
        // This moves the floating point division out of the hot loop
        layer.rescaleFactors[oc] = (inputScale * weightScale) / outputScale;

        // 4. Quantize Weights
        for (uint32_t i = 0; i < filterSize; ++i) {
            float v = fp32Weights[oc * filterSize + i] / weightScale;
            layer.weights[oc * filterSize + i] = static_cast<int8_t>(std::clamp(std::round(v), -127.0f, 127.0f));
        }
    }
    return layer;
}

#version 450
#extension GL_EXT_shader_dot_product : require
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require

layout(local_size_x = 16, local_size_y = 16) in;

layout(std430, set = 0, binding = 0) buffer Input  { uint  inData[]; };
layout(std430, set = 0, binding = 1) buffer Weight { uint  wData[];  };
layout(std430, set = 0, binding = 2) buffer Factors { float rescale[]; };
layout(std430, set = 0, binding = 3) buffer Output { uint  outData[]; };

void main() {
    uint outC = gl_GlobalInvocationID.z;
    ivec2 pos = ivec2(gl_GlobalInvocationID.xy);

    // 1. Initialize 32-bit Accumulator
    int sum = 0;

    // 2. The Vectorized Hot Loop
    // We process 4 channels at a time using a single instruction
    for(uint i = 0; i < (inputChannels / 4); ++i) {
        uint a = inData[idx_in + i];
        uint b = wData[idx_w + i];
        sum = dotProduct4x8bitPackedEXT(a, b, sum);
    }

    // 3. The Fixed-Point Rescale
    // We bring the high-precision sum back to the 8-bit range of the next layer
    float result = float(sum) * rescale[outC];

    // 4. Re-pack into the output buffer
    // Note: In a real shader, you would collect 4 results and write one 'uint'
    store_packed_result(pos, outC, int8_t(clamp(round(result), -128, 127)));
}