Trading Precision for Speed: Implementing Model Quantization and Format Conversion for Clinical Medical Imaging

How to Shrink Your Medical AI Model Without Losing Diagnostic Accuracy

The Introduction: When Every Millisecond and Megabyte Matter in Healthcare

The Problem:

In clinical environments, deploying deep learning models for medical imaging—brain tumors, strokes, cardiac anomalies—presents a unique constraint: the inference latency and memory footprint directly impact patient outcomes. A radiologist’s diagnostic workflow cannot tolerate 5-second model inference times when urgent conditions are at stake. Additionally, healthcare institutions operate on heterogeneous hardware: edge devices in clinics, CPU-only servers in remote facilities, and GPU-accelerated labs in major hospitals. Distributing a monolithic 500MB PyTorch model to all these environments is operationally untenable.

The VitalCheck API solves this by implementing model quantization (converting 32-bit floating-point weights to 8-bit integers) and format standardization (converting PyTorch models to ONNX—Open Neural Network Exchange). This approach achieves 4-6x model compression, 2-3x inference speedup, and universal hardware compatibility—all without sacrificing the diagnostic precision required in clinical settings.

The Solution:

We’ll analyze the quantization and deployment pipeline in app/ml/imaging.py and scripts/train_image_model.py, focusing on:

• Why int8 quantization works for medical imaging
• The ONNX conversion process and its advantages
• Building inference pipelines that validate quantized model accuracy
• Strategies for maintaining clinical-grade reliability with reduced precision

By the end of this tutorial, you’ll understand how to:

• Quantize PyTorch models while preserving diagnostic accuracy using calibration techniques
• Convert models to ONNX with custom operators and metadata for healthcare workflows
• Implement dual-format inference (both .pt and .onnx) with automatic fallback logic
• Validate quantized models against reference implementations using statistical methods
• Design inference pipelines that balance latency and accuracy in production

Essential Prerequisites: The Knowledge and Tools You Need

Knowledge Base:

Before diving in, you should be comfortable with:

• Deep Learning Fundamentals: Understanding neural network architecture, forward passes, and weight tensors (shape, dtype, gradients)
• PyTorch Model Development: Familiarity with torch.nn.Module, state_dict serialization, and model.eval() vs. model.train()
• Numerical Precision: Conceptual understanding of IEEE 754 floating-point representation and integer quantization (uniform vs. non-uniform)
• Healthcare ML Context: Awareness that medical models require higher reliability standards than consumer applications (FDA clearance, clinical validation)
• API Design: Basic FastAPI knowledge from the VitalCheck codebase (Pydantic schemas, dependency injection)

Environment Setup:

The VitalCheck API is built with these critical dependencies (from pyproject.toml):

Python 3.10+
PyTorch 2.0.1+ (for quantization-aware training and ONNX export)
ONNX 1.14.0+ (Open Neural Network Exchange runtime)
ONNXRuntime 1.15.0+ (cross-platform inference engine)
FastAPI 0.100.0+ (REST API framework)
Pydantic 2.0+ (data validation and serialization)
NumPy 1.24+ (numerical operations and calibration)

Critical Tools:

• PyTorch Quantization Tools: torch.quantization, torch.ao.quantization (advanced quantization aware training)
• ONNX Converter: torch.onnx.export() with opset version compatibility checking
• ONNXRuntime Inference: onnxruntime.InferenceSession for cross-platform execution
• Validation Framework: Custom accuracy metrics that respect medical imaging requirements (sensitivity > 95%, specificity > 90% for brain tumor classification)

The Pipeline as a Pharmaceutical Supply Chain: Architecture and Data Flow

Think of model quantization like a pharmaceutical supply chain optimization:

• Full-Precision Training = Producing medicine at high concentration (expensive, stable)
• Quantization = Diluting medicine to a clinically-effective concentration (cheaper, faster, same outcome)
• Calibration = Running trials to ensure the diluted version works as well as the original
• ONNX Export = Standardizing the diluted medicine’s format so any hospital (Intel, ARM, GPU) can use it
• Runtime Fallback = Having the original formula on hand if the diluted version shows unexpected reactions

Here’s the data flow architecture:

graph TD
    A["Raw Medical Images<br/>Brain MRI DICOM"] -->|Preprocessing| B["Normalized Tensors<br/>224x224x3 RGB"]
    B -->|Full-Precision Model<br/>32-bit float| C["Baseline Logits<br/>Confidence Scores"]
    
    D["Training Dataset<br/>Calibration Subset"] -->|Quantization-Aware Training| E["QAT Model<br/>Simulated int8"]
    E -->|ONNX Export| F["ONNX Format<br/>Quantized Weights"]
    
    C -->|Label Encoding| G["Reference Predictions<br/>4-class brain tumors"]
    F -->|ONNXRuntime Inference| H["Quantized Predictions<br/>int8 computation"]
    
    G -->|Comparison| I{"Accuracy ≥<br/>Baseline - 1%?"}
    H -->|Comparison| I
    
    I -->|✓ Approved| J["model_int8.onnx<br/>4.2x Compressed"]
    I -->|✗ Rejected| K["Fallback to<br/>model_fp32.onnx"]
    
    J -->|Runtime Selection| L["FastAPI Endpoint<br/>Choose best format"]
    K -->|Runtime Selection| L
    
    L -->|REST Request| M["Inference<br/>2-3x Faster"]
    M -->|JSON Response| N["Radiologist UI<br/>Tumor Class + Confidence"]

Key Flow Principles:

1. Dual-Format Storage: Maintain both quantized (int8) and full-precision (fp32) ONNX models
2. Calibration-Based Validation: Quantization quality is proven via statistical comparison, not just file size reduction
3. Runtime Optimization: The inference endpoint selects the optimal format based on available hardware
4. Fallback Safety: If quantized model degrades accuracy, seamlessly revert to full-precision

Building the Quantization Pipeline: A Step-by-Step Journey Through Precision Reduction

Stage 1: Understanding the Quantization Strategy

The Concept:

Quantization is the process of mapping 32-bit floating-point numbers to 8-bit integers. A single weight of 0.3847 (4 bytes) becomes 42 (1 byte). This works because:

• Neural network weights cluster around small ranges (typically -1 to +1)
• The relative magnitude between weights matters more than absolute precision
• Medical imaging benefits from 95%+ accuracy, not 99.9%

Code Block 1: Setting Up the Quantization Configuration

import torch
import torch.quantization as tq
from torch.ao.quantization import (
    get_default_qconfig_mapping,
    prepare_fx,
    convert_to_reference_real_quantized,
    QConfigMapping
)

# This configuration tells PyTorch:
# - Use static quantization (calibrate on fixed dataset, not per-batch)
# - Use symmetric quantization (range from -128 to +127)
# - Apply per-channel quantization for weights (more accurate than per-layer)
qconfig_mapping = (
    QConfigMapping()
    .set_global(tq.get_default_qconfig("fbgemm"))  # CPU-optimized quantization
)

# For medical models, we use a more conservative approach:
model.qconfig = tq.QConfig(
    activation=tq.HistogramObserver.with_args(reduce_range=False),
    weight=tq.PerChannelMinMaxObserver.with_args(dtype=torch.qint8)
)

print(f"✓ Quantization config ready. Weight precision: 8-bit integer per-channel")

Explanation:

• HistogramObserver calibrates activation ranges by analyzing the full distribution (not just min/max), crucial for preserving accuracy in medical networks
• PerChannelMinMaxObserver treats each output channel separately—this is critical because different filters in convolutional layers have different weight distributions
• reduce_range=False means we use the full [-128, 127] range instead of [-64, 63], important for healthcare models where we can’t afford to lose bit precision

Stage 2: Calibrating on Reference Data

The Logic:

Quantization requires a calibration dataset. The model runs on representative medical images, and PyTorch observes the range of activations and weights. For medical imaging:

• Use ~100-500 diverse images from the training set
• Ensure all tumor classes are represented (balanced calibration)
• This is different from model training—we’re just collecting statistics

Code Block 2: Building a Calibration Loop

# NOT training—just observing activation ranges
model.eval()

# Prepare model for quantization calibration
model_fp32 = model.to('cpu')  # Quantization happens on CPU
model_fp32.qconfig = torch.quantization.get_default_qconfig('fbgemm')
torch.quantization.prepare(model_fp32, inplace=True)

# Load representative brain MRI images
calibration_loader = torch.utils.data.DataLoader(
    calibration_dataset,  # ~400 images, all 4 tumor classes
    batch_size=16,
    shuffle=False,
    num_workers=0  # Important: Quantization is memory-sensitive
)

# Run calibration loop—forward pass only, no backprop
with torch.no_grad():  # Don't compute gradients (not needed)
    for images, _ in tqdm(calibration_loader, desc="Calibrating quantization"):
        # shapes: images = [16, 3, 224, 224]
        output = model_fp32(images)  # PyTorch observes activations internally
        # The histogram observer now knows: "max activation in this layer was 2.47"

print("✓ Calibration complete. Activation ranges recorded for all 47 layers.")

Why This Matters for Medical Imaging:

Without proper calibration, a quantized model might:

• Use a [0, 1] activation range when actual values reach 5.2 (overflow → clipping → accuracy loss)
• Over-allocate range for a layer that only uses 10% of it (underutilized precision)

Healthcare applications cannot tolerate this. A missed tumor due to poor quantization is a regulatory and clinical failure.

Stage 3: Converting to Quantized Model

Code Block 3: The Quantization Conversion

# Post-training quantization (PTQ)—convert the calibrated model
model_int8 = torch.quantization.convert(model_fp32, inplace=False)

# Verify the model structure changed:
# Before: Conv2d weights dtype = torch.float32
# After:  Conv2d weights dtype = torch.qint8

print("Model quantized!")
print(f"Original model size: {get_size_mb(model_fp32):.2f} MB")  # ~123 MB
print(f"Quantized model size: {get_size_mb(model_int8):.2f} MB")  # ~30 MB
# Compression ratio: ~4.1x

# Helper function for size calculation:
def get_size_mb(model):
    param_size = 0
    for param in model.parameters():
        param_size += param.nelement() * param.element_size()
    return param_size / (1024 ** 2)

Critical Detail:

The model is still a PyTorch model at this point. The weights are now int8 internally, but when you access state_dict(), PyTorch stores them efficiently. However, PyTorch’s int8 inference is slower on GPUs. That’s why we convert to ONNX.

Stage 4: Exporting to ONNX Format

The Why:

ONNX is a hardware-agnostic format. PyTorch is optimized for NVIDIA GPUs. But your medical clinic’s radiology lab has an Intel server and an ARM edge device. ONNX + ONNXRuntime allows a single model file to run on:

• Intel/AMD CPUs (optimized via GEMM kernels)
• ARM processors (optimized for mobile/edge)
• NVIDIA GPUs (via CUDA provider)
• Apple Silicon (via CoreML conversion)

Code Block 4: ONNX Export with Validation

import onnx
import onnxruntime as ort
from onnx import TensorProto

# Define input specification
# Brain tumor model expects: [N, 3, 224, 224] images
dummy_input = torch.randn(1, 3, 224, 224, dtype=torch.float32)

# Export quantized model to ONNX
onnx_path_int8 = "models/brain_tumor_4class_int8.onnx"
torch.onnx.export(
    model_int8,
    dummy_input,
    onnx_path_int8,
    input_names=["image"],
    output_names=["logits"],
    opset_version=14,  # Version 14 has robust int8 support
    do_constant_folding=True,  # Optimize constant operations
    verbose=False,
    dynamic_axes={'image': {0: 'batch_size'}}  # Allow variable batch sizes
)

# Also export full-precision version for comparison/fallback
onnx_path_fp32 = "models/brain_tumor_4class.onnx"
torch.onnx.export(
    model_fp32, dummy_input, onnx_path_fp32,
    input_names=["image"], output_names=["logits"],
    opset_version=14, dynamic_axes={'image': {0: 'batch_size'}}
)

# Validate ONNX files are well-formed
onnx_model_int8 = onnx.load(onnx_path_int8)
onnx_model_fp32 = onnx.load(onnx_path_fp32)
onnx.checker.check_model(onnx_model_int8)
onnx.checker.check_model(onnx_model_fp32)

print("✓ Both ONNX models exported and structurally validated")

What ONNX Export Does:

1. Traces the model execution (runs one forward pass)
2. Records every operation (Conv2d → BatchNorm → ReLU → etc.)
3. Serializes weights and operation graph to .onnx protobuf file
4. Supports quantized int8 operations (opset 14+)
5. Removes PyTorch-specific code (no Python dependencies at inference time)

Stage 5: Cross-Validation—Proving Quantization Didn’t Hurt Accuracy

Code Block 5: Quantization Accuracy Validation

import numpy as np
from sklearn.metrics import accuracy_score, precision_recall_fscore_support

# Load both models via ONNXRuntime
session_int8 = ort.InferenceSession(
    onnx_path_int8,
    providers=['CPUExecutionProvider']  # Use CPU for deterministic results
)
session_fp32 = ort.InferenceSession(
    onnx_path_fp32,
    providers=['CPUExecutionProvider']
)

# Validation on held-out test set (NOT calibration data)
validation_loader = torch.utils.data.DataLoader(
    validation_dataset,  # ~500 images
    batch_size=1,
    shuffle=False
)

predictions_int8 = []
predictions_fp32 = []
ground_truth = []

with torch.no_grad():
    for images, labels in tqdm(validation_loader, desc="Validating quantization"):
        # Convert torch tensor to numpy for ONNX inference
        image_np = images.numpy().astype(np.float32)  # shape: [1, 3, 224, 224]
        
        # Inference on int8 model
        logits_int8 = session_int8.run(
            None,
            {"image": image_np}
        )[0]  # Output is list of arrays, take first (logits)
        
        # Inference on fp32 model
        logits_fp32 = session_fp32.run(
            None,
            {"image": image_np}
        )[0]
        
        # Argmax to get class predictions
        pred_int8 = np.argmax(logits_int8, axis=1)[0]
        pred_fp32 = np.argmax(logits_fp32, axis=1)[0]
        
        predictions_int8.append(pred_int8)
        predictions_fp32.append(pred_fp32)
        ground_truth.append(labels.numpy()[0])

# Compute metrics
accuracy_int8 = accuracy_score(ground_truth, predictions_int8)
accuracy_fp32 = accuracy_score(ground_truth, predictions_fp32)
precision_int8, recall_int8, f1_int8, _ = precision_recall_fscore_support(
    ground_truth, predictions_int8, average='weighted'
)

print(f"FP32 Accuracy: {accuracy_fp32:.4f}")
print(f"Int8 Accuracy: {accuracy_int8:.4f}")
print(f"Accuracy Drop: {(accuracy_fp32 - accuracy_int8):.4f} ({(accuracy_fp32 - accuracy_int8)*100:.2f}%)")
print(f"\nInt8 Model Metrics:")
print(f"  Precision: {precision_int8:.4f}")
print(f"  Recall: {recall_int8:.4f}")
print(f"  F1-Score: {f1_int8:.4f}")

# Acceptance criterion: Max 1% accuracy drop for medical models
if accuracy_fp32 - accuracy_int8 <= 0.01:
    print("✓ APPROVED: Quantized model meets clinical accuracy requirements")
else:
    print("✗ REJECTED: Accuracy drop exceeds 1%. Use full-precision model only.")

Medical-Grade Validation:

• We use precision/recall/F1, not just accuracy
• For 4-class brain tumor classification:
  • High Recall (≥95%) = Don’t miss tumors (false negatives are dangerous)
  • High Precision (≥90%) = Minimize false alarms (unnecessary biopsies)
• We test on held-out validation data, never calibration data (prevents overfitting to calibration)

The Quantization Mechanics: Why 8-Bit Integers Work (And When They Don’t)

Understanding Linear Quantization

The Math:

Quantization maps floating-point range [min, max] to integer range [-128, 127]:

quantized_value = round(
    (fp32_value - zero_point) * scale
)

where:
  scale = 127 / (max_fp32 - min_fp32)
  zero_point = (min_fp32 * scale)

Concrete Example from Brain Tumor Model:

Layer: “Conv2d_1.weight”

• Original range: [-0.285, +0.412] (float32)
• Quantization range: 0.412 - (-0.285) = 0.697
• Scale factor: 127 / 0.697 = 182.07
• A weight of 0.123 becomes: round((0.123 - (-0.115)) * 182.07) = round(43.4) = 43
• Dequantized back: 43 / 182.07 + (-0.115) = 0.221 - 0.115 = 0.106
• Error: 0.123 - 0.106 = 0.017 (~1.4% relative error)

Why This Works for Medical Imaging:

1. Relative Magnitude Preservation: A weight that was 2x another weight is still ~2x after quantization
2. Activation Preservation: The model’s decision boundaries (where it transitions from “Glioma” to “Meningioma”) are determined by activation patterns, which tolerate small quantization noise
3. Statistical Redundancy: Neural networks are trained with dropout and regularization—they inherently handle small perturbations

When This Fails:

🔴 Danger: Quantization can severely degrade accuracy in:

• Low-Bit Quantization (binary/ternary—1-2 bits per weight) because the quantization grid becomes too coarse
• Models trained with batch normalization that fold into weights (BN moves to lower precision first)
• Models with extreme weight distributions (some weights -0.001 to +0.001, others -100 to +100)

ONNXRuntime’s Inference Optimization

Code Block 6: Why ONNX Int8 Inference is Faster

import time

# Benchmark: PyTorch float32 vs ONNX int8

# Setup
test_image = torch.randn(1, 3, 224, 224)
iterations = 100

# PyTorch int8 inference (on GPU, this is slow)
model_int8_pytorch = model_int8.to('cuda')
torch.cuda.synchronize()
start = time.time()
with torch.no_grad():
    for _ in range(iterations):
        _ = model_int8_pytorch(test_image.to('cuda'))
torch.cuda.synchronize()
pytorch_time = time.time() - start

# ONNX int8 inference (CPU, optimized)
session_int8_cpu = ort.InferenceSession(
    onnx_path_int8,
    providers=['CPUExecutionProvider']
)
test_image_np = test_image.numpy().astype(np.float32)
start = time.time()
for _ in range(iterations):
    _ = session_int8_cpu.run(None, {"image": test_image_np})
onnx_time = time.time() - start

print(f"PyTorch int8 (GPU): {pytorch_time/iterations*1000:.2f} ms/inference")
print(f"ONNX int8 (CPU):    {onnx_time/iterations*1000:.2f} ms/inference")
# Typical result: PyTorch 45ms, ONNX 12ms (3.75x faster!)

Why the Speedup?

1. CPU Vectorization: ONNXRuntime uses GEMM (General Matrix Multiply) kernels optimized for int8 on AVX2/AVX-512 CPU instructions
2. No Runtime Compilation: ONNX model is static—no Python GIL, no dynamic graph construction
3. Memory Bandwidth: int8 uses 4x less memory bandwidth than float32—critical for inference in production
4. Operator Fusion: ONNXRuntime fuses conv → relu → pooling into single kernel, eliminating intermediate buffer allocations

🔵 Deep Dive: Modern CPUs have dedicated int8 execution units. An Intel CPU’s VNNI (Vector Neural Network Instructions) can multiply four pairs of int8 values in a single clock cycle, vs. float32 requiring full IEEE arithmetic.

The Danger Zone: When Quantization Fails in Clinical Settings

Pitfall 1: Quantization Mismatch Between Training and Inference

The Problem:

# ❌ WRONG: Train with full precision, then quantize
model = train_model(...)  # PyTorch, float32
model_int8 = torch.quantization.convert(model)  # Post-training quantization (PTQ)

# The model never "learned" to work with quantized activations during training.
# When you deploy int8, activations follow different distributions than during training.
# Result: 3-5% accuracy drop for quantization-sensitive architectures

The Solution:

# ✓ CORRECT: Quantization-Aware Training (QAT)
model = BrainTumorClassifier()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model = torch.quantization.prepare_qat(model)  # Insert fake quantization ops

# Train normally—but now the model learns with quantization in the loop
optimizer = torch.optim.Adam(model.parameters())
for epoch in range(num_epochs):
    for images, labels in train_loader:
        logits = model(images)
        loss = criterion(logits, labels)
        loss.backward()
        optimizer.step()

# After training, calibration on val data, then convert
model = torch.quantization.convert(model.eval())

Code Block 7: Detecting the Mismatch

# This test catches poor quantization early:
def validate_quantization_mismatch(model_fp32, model_int8, test_loader):
    """Check if quantized model has unexpected accuracy drops."""
    
    with torch.no_grad():
        fp32_correct = 0
        int8_correct = 0
        mismatch_count = 0
        
        for images, labels in test_loader:
            logits_fp32 = model_fp32(images)
            logits_int8 = model_int8(images)
            
            preds_fp32 = torch.argmax(logits_fp32, dim=1)
            preds_int8 = torch.argmax(logits_int8, dim=1)
            
            fp32_correct += (preds_fp32 == labels).sum().item()
            int8_correct += (preds_int8 == labels).sum().item()
            
            # Count where models disagree
            mismatch = (preds_fp32 != preds_int8).sum().item()
            mismatch_count += mismatch
            
            # 🔴 If they disagree on >= 5% of samples, quantization is poor
            if mismatch / len(labels) > 0.05:
                print(f"⚠️ High model disagreement: {mismatch/len(labels)*100:.1f}%")
    
    print(f"Agreement rate: {(test_count - mismatch_count)/test_count*100:.1f}%")

Pitfall 2: Calibration on Biased Data

The Problem:

# ❌ WRONG: Calibrate only on "easy" images
calibration_dataset = [
    img for img in all_images 
    if brain_tumor_is_obvious and image_quality_is_high
]
# Result: Quantization optimizes for high-contrast tumors.
# Real-world MRIs with artifacts / subtle tumors → poor inference!

The Solution:

# ✓ CORRECT: Calibrate on diverse, representative data
def build_balanced_calibration_set(dataset, tumor_classes=4, samples_per_class=100):
    """Ensure all classes and difficulty levels are represented."""
    
    calibration = {
        'obvious': [],      # High SNR, clear tumor boundaries
        'subtle': [],       # Low SNR, unclear boundaries
        'artifact': [],     # Motion artifacts, scanner noise
        'normal': []        # Healthy brain scans (negative control)
    }
    
    for img, label, metadata in dataset:
        difficulty = assess_image_quality(img, metadata)  # Your quality metric
        
        if len(calibration[difficulty]) < samples_per_class:
            calibration[difficulty].append((img, label))
    
    # Flatten and verify distribution
    calibration_set = []
    for difficulty_level in calibration.values():
        calibration_set.extend(difficulty_level)
    
    assert len(calibration_set) == 4 * samples_per_class
    print(f"✓ Calibration set balanced: {len(calibration_set)} images across 4 difficulty levels")
    
    return calibration_set

Pitfall 3: Fallback Logic Failures in Distributed Inference

The Problem:

In a hospital network, you have:

• Edge device (clinic exam room): ARM CPU, 2 GB RAM → only int8 fits
• Server (radiology dept): Intel x86, 64 GB RAM → can handle fp32
• Cloud backup: GPU → expects float32

If all devices use hardcoded quantization settings:

# ❌ Fragile: Assumes int8 works everywhere
session = ort.InferenceSession("models/brain_tumor_int8.onnx")
logits = session.run(None, {"image": img_np})

If the edge device goes offline, the server doesn’t know which model file to use.

The Solution:

Code Block 8: Adaptive Model Selection with Validation

import onnxruntime as ort
from typing import Tuple, Optional

class AdaptiveModelSelector:
    """Chooses best model format based on hardware and accuracy validation."""
    
    def __init__(self, model_int8_path: str, model_fp32_path: str):
        self.model_int8_path = model_int8_path
        self.model_fp32_path = model_fp32_path
        self.session_int8 = None
        self.session_fp32 = None
        self.selected_model = None
        
        # Metadata: accuracy metrics for fallback decisions
        self.accuracy_int8 = 0.9487  # From validation
        self.accuracy_fp32 = 0.9534
        self.accuracy_margin = 0.0047  # 0.47% drop
    
    def initialize(self, prefer_quantized: bool = True) -> str:
        """Attempt to load both models, select best available."""
        
        # Try quantized model first
        try:
            self.session_int8 = ort.InferenceSession(
                self.model_int8_path,
                providers=['CPUExecutionProvider', 'CUDAExecutionProvider']
            )
            print(f"✓ Loaded int8 model. Execution provider: {self.session_int8.get_providers()}")
            self.selected_model = 'int8'
            
        except Exception as e:
            print(f"⚠️ Failed to load int8 model: {e}")
            self.session_int8 = None
        
        # Try full-precision model as fallback
        try:
            self.session_fp32 = ort.InferenceSession(
                self.model_fp32_path,
                providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
            )
            print(f"✓ Loaded fp32 model. Execution provider: {self.session_fp32.get_providers()}")
            
            if self.selected_model is None:
                self.selected_model = 'fp32'
                
        except Exception as e:
            print(f"🔴 CRITICAL: Failed to load fp32 model: {e}")
            raise RuntimeError("No inference models available!")
        
        return self.selected_model
    
    def infer(self, image_np: np.ndarray, validate: bool = False) -> Tuple[np.ndarray, str]:
        """Run inference with automatic fallback if selected model fails."""
        
        if self.selected_model is None:
            raise RuntimeError("Call initialize() before inference!")
        
        try:
            # Try selected model first
            if self.selected_model == 'int8' and self.session_int8:
                logits = self.session_int8.run(None, {"image": image_np})[0]
                
                # Optional: Cross-validate against fp32 in production
                if validate and self.session_fp32:
                    logits_fp32 = self.session_fp32.run(None, {"image": image_np})[0]
                    pred_int8 = np.argmax(logits[0])
                    pred_fp32 = np.argmax(logits_fp32[0])
                    
                    if pred_int8 != pred_fp32:
                        print(f"⚠️ Model disagreement detected! Int8: {pred_int8}, FP32: {pred_fp32}")
                        print("    Falling back to full-precision inference...")
                        logits = logits_fp32
                        self.selected_model = 'fp32'
                
                return logits, 'int8'
            
            elif self.selected_model == 'fp32' and self.session_fp32:
                logits = self.session_fp32.run(None, {"image": image_np})[0]
                return logits, 'fp32'
        
        except ort.NotImplementedError as e:
            # Model op not supported on this hardware
            print(f"⚠️ Hardware doesn't support selected model: {e}")
        
        except Exception as e:
            print(f"⚠️ Inference error with {self.selected_model}: {e}")
        
        # Fallback: Try the other model
        print("Attempting fallback model...")
        try:
            if self.selected_model == 'int8' and self.session_fp32:
                logits = self.session_fp32.run(None, {"image": image_np})[0]
                self.selected_model = 'fp32'  # Switch permanently
                return logits, 'fp32 (fallback)'
            
            elif self.selected_model == 'fp32' and self.session_int8:
                logits = self.session_int8.run(None, {"image": image_np})[0]
                self.selected_model = 'int8'  # Switch permanently
                return logits, 'int8 (fallback)'
        
        except Exception as e:
            print(f"🔴 CRITICAL: Both models failed! {e}")
            raise RuntimeError("Complete inference failure—no models available")
        
        raise RuntimeError("No fallback model available")

# Usage in FastAPI endpoint:
selector = AdaptiveModelSelector("models/brain_tumor_int8.onnx", "models/brain_tumor_fp32.onnx")
selected = selector.initialize(prefer_quantized=True)
print(f"Using {selected} model for inference")

# In request handler:
@app.post("/predict")
async def predict_brain_tumor(file: UploadFile):
    image = preprocess_mri(file)
    logits, model_used = selector.infer(image, validate=True)  # validate=True in production
    prediction = torch.argmax(torch.tensor(logits)).item()
    return {
        "tumor_class": CLASS_NAMES[prediction],
        "model_used": model_used,
        "confidence": float(torch.softmax(torch.tensor(logits), dim=1)[0, prediction])
    }

Why This Matters:

In a clinic, a hardware failure (GPU dies, runs out of memory) cannot cause diagnostic failure. The fallback logic ensures that even if the optimized int8 model has issues, the fp32 version is ready.

Pitfall 4: Batch Size Mismatches in Dynamic Axes

The Problem:

# Your ONNX model trained with batch_size=16
# But a clinic's edge device sends batch_size=1 images one at a time
# Quantization may behave differently across batch sizes!

# Reason: Batch normalization and activation statistics change with batch size

Detection and Fix:

# Code Block 9: Test dynamic batch sizes
def test_batch_size_robustness(session, test_image_np):
    """Verify model predictions are consistent across batch sizes."""
    
    single_batch = test_image_np  # shape: [1, 3, 224, 224]
    dual_batch = np.concatenate([test_image_np, test_image_np], axis=0)  # [2, 3, 224, 224]
    quad_batch = np.tile(test_image_np, (4, 1, 1, 1))  # [4, 3, 224, 224]
    
    logits_1 = session.run(None, {"image": single_batch})[0]
    logits_2 = session.run(None, {"image": dual_batch})[0]
    logits_4 = session.run(None, {"image": quad_batch})[0]
    
    # Predictions should be identical regardless of batch size
    pred_1 = np.argmax(logits_1[0])  # From single batch
    pred_2_first = np.argmax(logits_2[0])  # First image in dual batch
    pred_4_first = np.argmax(logits_4[0])  # First image in quad batch
    
    assert pred_1 == pred_2_first == pred_4_first, \
        f"Batch size affects predictions! {pred_1} vs {pred_2_first} vs {pred_4_first}"
    
    print("✓ Batch size robustness verified")

Mastering Model Deployment: Your New Superpower

The Skill You’ve Acquired

You now understand:

1. Quantization as a Medical Tool: You can compress models 4-6x while maintaining clinical accuracy through rigorous calibration and validation

2. Format Optimization: You know why ONNX enables universal deployment and how ONNXRuntime provides 2-3x speedup over native PyTorch inference

3. Accuracy Validation: You can statistically prove that quantized models meet medical requirements (precision/recall > thresholds) using held-out test sets

4. Distributed Resilience: You can design fallback logic that ensures inference never fails in production, automatically switching between quantized and full-precision models

5. Edge Case Handling: You understand pitfalls (calibration bias, batch size effects, training-inference mismatch) and how to detect them before deployment

Practical Next Steps

For Your VitalCheck API:

1. Apply QAT to all models:

   python scripts/train_image_model.py --quantize --qat  # Enables QAT during training
   python scripts/train_tabular_models.py --quantize     # Quantize tabular models too

2. Validate on realistic data:

   python -m app.ml.imaging validate --model-int8 data/models/brain_tumor_int8.onnx

3. Deploy with fallback:

   • Update app/ml/imaging.py to use AdaptiveModelSelector from Code Block 8
   • Test in staging with network failures, memory pressure, hardware mismatches

4. Monitor in production:

   • Log which model version was used (model_used field in response)
   • Alert if fallback rate exceeds 5% (indicates hardware issues)

One Final Thing: The Medical Mindset

In consumer ML, 95% vs 94% accuracy is often acceptable. In clinical ML, this percentage point difference can mean:

• 1% of brain tumors missed
• 1% of stroke patients not getting timely intervention
• Regulatory violations and loss of FDA clearance

Every optimization technique—quantization, pruning, distillation—must be validated with this stakes in mind. Your quantization pipeline does this. Use it.