quantize_model.py

import tensorflow as tf
import numpy as np
import time
import h5py
import os

# ==========================================
# CONFIGURATION
# ==========================================
NUM_QUBITS = 5
DOWNSAMPLE_FACTOR = 40
M_PARAM = 16
ORIGINAL_LENGTH = 1000
TIME_STEPS = ORIGINAL_LENGTH // DOWNSAMPLE_FACTOR

MODEL_NAME = f"qubit_model_ds{DOWNSAMPLE_FACTOR}_m{M_PARAM}.h5"
TEST_PATH = f"/home/sandra/Qubit_5Channel_ds{DOWNSAMPLE_FACTOR}_test.h5"

np.random.seed(42)


# ==========================================
# 1. LOAD TEST DATA (MATCH TRAINING EXACTLY)
# ==========================================
def load_test_data(path):
    with h5py.File(path, "r") as hf:
        X = hf["X_test"][:]
        y = hf["y_test"][:]

    # --- preprocessing (must match training) ---
    X_temp = X.reshape(-1, ORIGINAL_LENGTH // 10, 5, 2)

    q_scales = np.max(np.abs(X_temp), axis=(0, 1, 3), keepdims=True)
    X_scaled = X_temp / (q_scales + 1e-12)

    X = X_scaled.reshape(-1, ORIGINAL_LENGTH // 10, 10)
    X = X.reshape(-1, TIME_STEPS, 1, 10)

    # labels → bit representation
    y_bits = np.array([
        [(label >> i) & 1 for i in range(NUM_QUBITS)]
        for label in y
    ])

    return X.astype(np.float32), y_bits


# ==========================================
# 2. LOAD MODEL + DATA
# ==========================================
print("[INFO] Loading model + data...")
model = tf.keras.models.load_model(MODEL_NAME)
X_test, y_test = load_test_data(TEST_PATH)


# ==========================================
# 3. REPRESENTATIVE DATASET
# ==========================================
def representative_data_gen():
    idx = np.random.choice(len(X_test), size=min(1000, len(X_test)), replace=False)
    for i in idx:
        yield [X_test[i:i+1].astype(np.float32)]


# ==========================================
# 4. FULL INT8 CONVERSION
# ==========================================
print("\n[INFO] Converting to FULL INT8...")

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen

converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]

converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8

tflite_model = converter.convert()

TFLITE_PATH = f"model_int8_full_ds{DOWNSAMPLE_FACTOR}_m{M_PARAM}.tflite"

with open(TFLITE_PATH, "wb") as f:
    f.write(tflite_model)

print("[INFO] FULL INT8 model saved.")


# ==========================================
# 5. FP32 BASELINE
# ==========================================
print("\n[INFO] Running FP32 baseline...")

start = time.time()
fp32_raw = model.predict(X_test, verbose=0)
fp32_preds = (fp32_raw > 0.5).astype(int)
lat_fp32 = (time.time() - start) / len(X_test)


# ==========================================
# 6. FULL INT8 INFERENCE
# ==========================================
def run_tflite_int8(model_path, data):
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()

    input_details = interpreter.get_input_details()[0]
    output_details = interpreter.get_output_details()[0]

    # 🔍 quantization params
    in_scale, in_zero = input_details["quantization"]
    out_scale, out_zero = output_details["quantization"]

    print("\n[DEBUG] Input quant:", input_details["dtype"], in_scale, in_zero)
    print("[DEBUG] Output quant:", output_details["dtype"], out_scale, out_zero)

    preds = []
    start = time.time()

    for i in range(len(data)):
        x = data[i:i+1].astype(np.float32)

        # 🔹 quantize input
        x_int8 = (x / in_scale + in_zero).round().astype(np.int8)

        interpreter.set_tensor(input_details["index"], x_int8)
        interpreter.invoke()

        # 🔹 get int8 output
        out_int8 = interpreter.get_tensor(output_details["index"])

        # 🔹 dequantize output
        out = (out_int8.astype(np.float32) - out_zero) * out_scale

        preds.append(out[0])

    latency = (time.time() - start) / len(data)

    preds = np.array(preds)
    preds = (preds > 0.5).astype(int)

    return preds, latency


print("\n[INFO] Running FULL INT8 inference...")
int8_preds, lat_int8 = run_tflite_int8(TFLITE_PATH, X_test)


# ==========================================
# 7. METRICS
# ==========================================
def get_metrics(p, y):
    q_accs = [np.mean(p[:, i] == y[:, i]) for i in range(NUM_QUBITS)]
    ga = np.mean(np.all(p == y, axis=1))
    fid = np.exp(np.mean(np.log(np.array(q_accs) + 1e-12)))
    return ga, fid, q_accs


ga_fp32, fid_fp32, q_fp32 = get_metrics(fp32_preds, y_test)
ga_int8, fid_int8, q_int8 = get_metrics(int8_preds, y_test)


# ==========================================
# 8. MODEL SIZE
# ==========================================
fp32_size = os.path.getsize(MODEL_NAME) / (1024 * 1024)
int8_size = os.path.getsize(TFLITE_PATH) / (1024 * 1024)


# ==========================================
# 9. REPORT
# ==========================================
print("\n" + "="*60)
print(f"{'Metric':<18} | {'FP32':<12} | {'INT8 FULL':<12}")
print("="*60)

print(f"{'Latency (ms)':<18} | {lat_fp32*1000:<12.4f} | {lat_int8*1000:<12.4f}")
print(f"{'Global Acc':<18} | {ga_fp32*100:<11.2f}% | {ga_int8*100:<11.2f}%")
print(f"{'Fidelity':<18} | {fid_fp32*100:<11.2f}% | {fid_int8*100:<11.2f}%")

print("\nQubit Accuracies:")
for i in range(NUM_QUBITS):
    print(f"Q{i:<2}               | {q_fp32[i]*100:<11.2f}% | {q_int8[i]*100:<11.2f}%")

fp32_size = os.path.getsize(MODEL_NAME) / 1024
int8_size = os.path.getsize(TFLITE_PATH) / 1024

print("\nMODEL SIZE:")
print(f"{'FP32 (.h5)':<18}: {fp32_size:.2f} KB")
print(f"{'INT8 (.tflite)':<18}: {int8_size:.2f} KB")

speed_up= lat_fp32 / lat_int8 if lat_int8 > 0 else float('inf')
size_reduction = fp32_size / int8_size if int8_size > 0 else float('inf')
print(f"\nSpeed-up: {speed_up:.2f}x")
print(f"Size Reduction: {size_reduction:.2f}x")

print("="*60)