def main():
"""Convert a trained Keras model to quantized TFLite and validate."""
args = get_args()
# Building a backbone and reusing one are opposite intents; asking for both
# would silently discard whichever ran second.
if args.split_head and args.backbone_path:
raise SystemExit(
"--split_head builds a new backbone and --backbone_path reuses an existing one; pass only one."
)
if args.backbone_path and not os.path.isfile(args.backbone_path):
raise FileNotFoundError(f"Backbone model not found: {args.backbone_path}")
# Resolve config path
if not args.model_config:
args.model_config = os.path.splitext(args.checkpoint_path)[0] + "_model_config.json"
if not os.path.isfile(args.model_config):
raise FileNotFoundError(f"Model config JSON not found: {args.model_config}")
cfg = ModelConfig.load(args.model_config).to_dict()
# Load model
model = load_keras_model(args.checkpoint_path)
print(f"Loaded model from {args.checkpoint_path}")
# Build representative dataset generator
data_manifests: dict[str, dict] = {}
if os.path.isdir(args.data_path_train):
configured_classes = cfg.get("class_names") or None
file_paths, classes = load_file_paths_from_directory(args.data_path_train, classes=configured_classes)
if not file_paths:
raise ValueError("No training audio found for the classes in the model config.")
class_count = len({os.path.basename(os.path.dirname(path)) for path in file_paths})
stratified_paths = stratified_sample_paths(file_paths, args.num_samples, seed=42)
if len(stratified_paths) != args.num_samples:
raise ValueError(
f"Requested {args.num_samples} calibration paths but only {len(stratified_paths)} are available."
)
print(f"Representative dataset: {len(stratified_paths)} stratified samples from {class_count} folders.")
data_manifests["calibration"] = _manifest_record(stratified_paths, args.data_path_train)
def rep_data_gen(num_samples: int | None = None):
count = len(stratified_paths) if num_samples is None else min(num_samples, len(stratified_paths))
return representative_data_gen(stratified_paths, cfg, num_samples=count)
# Calibration and validation must be disjoint; overlap makes parity
# reports optimistic and invalidates a release gate.
val_paths_subset = stratified_sample_paths(
file_paths,
args.validate_samples,
seed=43,
exclude=set(stratified_paths),
)
if not val_paths_subset:
raise ValueError("No files remain for disjoint quantization validation.")
if len(val_paths_subset) != args.validate_samples:
raise ValueError(
f"Requested {args.validate_samples} validation paths but only {len(val_paths_subset)} are available."
)
print(f"Validation dataset: {len(val_paths_subset)} disjoint stratified samples.")
data_manifests["validation"] = _manifest_record(val_paths_subset, args.data_path_train)
def rep_data_gen_val():
return representative_data_gen(val_paths_subset, cfg, num_samples=len(val_paths_subset))
else:
print("No training data directory provided; generating random representative dataset.")
def rep_data_gen(num_samples: int | None = None):
count = args.num_samples if num_samples is None else num_samples
sr = int(cfg["sample_rate"])
cd = cfg["chunk_duration"]
T = int(sr * cd)
spec_width = int(cfg["spec_width"])
n_fft = int(cfg["fft_length"])
frontend = normalize_frontend_name(cfg["audio_frontend"])
num_mels = int(cfg["num_mels"])
fft_bins = hybrid_fft_bins(n_fft)
for _ in tqdm(range(count), desc="Random samples", unit="sample"):
if frontend == "librosa":
yield [np.random.rand(1, num_mels, spec_width, 1).astype(np.float32)]
elif frontend == "hybrid":
yield [np.random.rand(1, fft_bins, spec_width, 1).astype(np.float32)]
else:
yield [np.random.randn(1, T, 1).astype(np.float32)]
def rep_data_gen_val():
return rep_data_gen(num_samples=args.validate_samples)
# Output path
if not args.output_path:
args.output_path = os.path.splitext(args.checkpoint_path)[0] + "_quantized.tflite"
# Conversion is staged beside the destination and promoted atomically only
# after every numerical quality gate passes. A failed conversion must
# never leave a release-looking .tflite artifact behind.
output_dir = os.path.dirname(os.path.abspath(args.output_path))
os.makedirs(output_dir, exist_ok=True)
with tempfile.NamedTemporaryFile(
prefix=".quantizing-", suffix=".tflite", dir=output_dir, delete=False
) as tmp_handle:
tmp_path = tmp_handle.name
report: dict = {
"output_path": args.output_path,
"quantization": args.quantization,
"per_tensor": args.per_tensor,
"quality_gate_passed": False,
"data_manifests": data_manifests,
}
try:
convert_to_tflite(model, rep_data_gen, tmp_path, quantization=args.quantization, per_tensor=args.per_tensor)
n_runs = max(1, args.batch_validate) if args.batch_validate > 0 else 1
all_metrics: list[dict] = []
for run_idx in range(n_runs):
if n_runs > 1:
print(f"\n--- Validation run {run_idx + 1}/{n_runs} ---")
val_metrics = validate_models(model, tmp_path, rep_data_gen_val)
all_metrics.append(val_metrics)
# Aggregate metrics across runs. Input manifests are deterministic, so
# repeated runs measure runtime repeatability rather than resampling.
if n_runs > 1:
print(f"\n--- Batch validation summary ({n_runs} runs) ---")
for key in ["cosine_mean", "cosine_p05", "mse_mean", "mae_mean", "pearson_mean"]:
vals = [m[key] for m in all_metrics]
worst = min(vals) if "cosine" in key or "pearson" in key else max(vals)
mean = np.mean(vals)
print(f" {key}: mean={mean:.6f} worst={worst:.6f}")
report["batch_validation"] = {"n_runs": n_runs, "all_metrics": all_metrics}
val_metrics = dict(all_metrics[0])
val_metrics["cosine_mean"] = min(m["cosine_mean"] for m in all_metrics)
val_metrics["cosine_p05"] = min(m["cosine_p05"] for m in all_metrics)
else:
val_metrics = all_metrics[0]
report["validation"] = val_metrics
failures = []
cos_mean = val_metrics["cosine_mean"]
cos_p05 = val_metrics["cosine_p05"]
if args.min_cosine_sim > 0 and cos_mean < args.min_cosine_sim:
failures.append(f"mean cosine {cos_mean:.6f} < {args.min_cosine_sim:.4f}")
if args.min_cosine_p05 > 0 and cos_p05 < args.min_cosine_p05:
failures.append(f"p05 cosine {cos_p05:.6f} < {args.min_cosine_p05:.4f}")
if failures:
report["quality_gate_failures"] = failures
_write_report(args.report_json, report)
raise RuntimeError("Quantization quality check failed: " + "; ".join(failures))
report["quality_gate_passed"] = True
os.replace(tmp_path, args.output_path)
tmp_path = ""
print(f"TFLite model validated and saved to {args.output_path}")
# Save validation data
# Save labels only for a model that passed the gate.
if cfg.get("class_names"):
labels_path = os.path.splitext(args.output_path)[0] + "_labels.txt"
with open(labels_path, "w", encoding="utf-8") as handle:
handle.writelines(f"{name}\n" for name in cfg["class_names"])
print(f"Labels saved to {labels_path}")
validation_batches = [sample[0] for sample in rep_data_gen_val()]
validation_data = np.concatenate(validation_batches, axis=0)
if validation_data.shape[0] > 25:
validation_data = pick_random_samples(validation_data, 25)
val_path = os.path.splitext(args.output_path)[0] + "_validation_data.npz"
np.savez_compressed(val_path, data=validation_data)
print(f"Validation data saved to {val_path}")
# Backbone / classifier split. Opt-in: the firmware still runs a single
# network, so the pair is produced only when it is asked for. The
# monolithic model above is converted either way, which is what the
# board flashes; the split artifacts are what travels over the air.
if args.split_head:
report["split"] = _convert_split_head(model, cfg, args, rep_data_gen, rep_data_gen_val)
elif args.backbone_path:
report["split"] = _convert_head_only(model, cfg, args, rep_data_gen, rep_data_gen_val)
# ONNX export
if args.export_onnx:
onnx_path = os.path.splitext(args.output_path)[0] + ".onnx"
report["onnx_validation"] = _export_and_validate_onnx(model, onnx_path, rep_data_gen_val)
report["onnx_path"] = onnx_path
print(f"ONNX model validated and saved to {onnx_path}")
report["model_size_bytes"] = os.path.getsize(args.output_path)
report["keras_size_bytes"] = os.path.getsize(args.checkpoint_path)
report["compression_ratio"] = report["keras_size_bytes"] / max(report["model_size_bytes"], 1)
report["config"] = cfg
_write_report(args.report_json, report)
finally:
if tmp_path and os.path.exists(tmp_path):
os.unlink(tmp_path)