From 090be14a6a89797b5639b21d1cde9eb243344b79 Mon Sep 17 00:00:00 2001 From: slalom Date: Mon, 1 Jun 2026 16:53:45 -0400 Subject: [PATCH] add script to test steps in ai-hub --- examples/ai-hub/README.md | 117 ++++++++++++++++++++++ examples/ai-hub/prepare_inputs.py | 75 ++++++++++++++ examples/ai-hub/run_ai_hub.sh | 156 ++++++++++++++++++++++++++++++ examples/training/source/train.py | 4 - 4 files changed, 348 insertions(+), 4 deletions(-) create mode 100644 examples/ai-hub/README.md create mode 100755 examples/ai-hub/prepare_inputs.py create mode 100755 examples/ai-hub/run_ai_hub.sh diff --git a/examples/ai-hub/README.md b/examples/ai-hub/README.md new file mode 100644 index 0000000..947598f --- /dev/null +++ b/examples/ai-hub/README.md @@ -0,0 +1,117 @@ +# Qualcomm AI Hub Example + +This example takes the ONNX model produced by the SageMaker training example and runs the Qualcomm AI Hub upload workflow: + +1. Quantize +2. Compile +3. Validate +4. Profile +5. Download the compiled artifact + +## Prerequisites + +Run the training example first and wait for it to complete: + +```bash +bash examples/training/run_training.sh --config config.yaml --wait +``` + +If the dataset is already uploaded to S3, use: + +```bash +bash examples/training/run_training.sh --config config.yaml --skip-upload --wait +``` + +The training artifact must contain a static-shape `model.onnx`. The training example exports an input named `input` with shape `1x3x160x160`. + +Your `config.yaml` must include AI Hub settings: + +```yaml +aihub: + device: Samsung Galaxy S25 (Family) + target_runtime: tflite + input_specs: + input: [[1, 3, 160, 160], float32] + output_dir: build/qai-hub +``` + +You also need local Qualcomm AI Hub SDK authentication configured. + +## Prepare Inputs + +AI Hub does not consume the raw JPG training images directly. It needs NumPy tensors that match the ONNX model input shape and preprocessing. + +Generate calibration and validation inputs: + +```bash +uv run python examples/ai-hub/prepare_inputs.py +``` + +This writes: + +```text +examples/training/data/aihub_calibration/*.npy +examples/training/data/inputs.npz +``` + +The script applies the same image preprocessing used by the training example: + +- resize to `160x160` +- convert to channel-first `1x3x160x160` +- normalize with ImageNet mean and standard deviation + +Useful options: + +```bash +uv run python examples/ai-hub/prepare_inputs.py \ + --dataset-dir examples/training/data/flower_photos_sagemaker \ + --calibration-dir examples/training/data/aihub_calibration \ + --input-file examples/training/data/inputs.npz \ + --samples 16 +``` + +## Run AI Hub + +After training completes and inputs are prepared: + +```bash +bash examples/ai-hub/run_ai_hub.sh --config config.yaml +``` + +By default, the script uses the last SageMaker training job recorded in `.qc-cli.json`. It downloads that job's `model.tar.gz`, extracts `model.onnx`, runs the AI Hub workflow, and downloads the compiled artifact. + +To use a specific training job: + +```bash +bash examples/ai-hub/run_ai_hub.sh \ + --config config.yaml \ + --from-job qc-cli-YYYYMMDD-HHMMSS +``` + +To resume from a later Workbench step: + +```bash +bash examples/ai-hub/run_ai_hub.sh \ + --config config.yaml \ + --from-step validate +``` + +To skip downloading the compiled artifact: + +```bash +bash examples/ai-hub/run_ai_hub.sh \ + --config config.yaml \ + --skip-download +``` + +## Troubleshooting + +If AI Hub reports dynamic input shapes, rerun training with the current training source. AI Hub quantization requires the exported ONNX model to use static input shapes. + +If `run_ai_hub.sh` reports missing calibration or input files, run: + +```bash +uv run python examples/ai-hub/prepare_inputs.py +``` + +If validation fails with a missing input name, make sure `config.yaml` and the generated `.npz` both use `input` as the input name. diff --git a/examples/ai-hub/prepare_inputs.py b/examples/ai-hub/prepare_inputs.py new file mode 100755 index 0000000..5f4c6e6 --- /dev/null +++ b/examples/ai-hub/prepare_inputs.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +"""Prepare Qualcomm AI Hub calibration and validation inputs for the training example.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +from PIL import Image + + +IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--dataset-dir", + type=Path, + default=Path("examples/training/data/flower_photos_sagemaker"), + help="ImageFolder-style dataset used for training.", + ) + parser.add_argument( + "--calibration-dir", + type=Path, + default=Path("examples/training/data/aihub_calibration"), + help="Directory where .npy calibration samples will be written.", + ) + parser.add_argument( + "--input-file", + type=Path, + default=Path("examples/training/data/inputs.npz"), + help="Validation .npz input file for qc-cli ai-hub validate.", + ) + parser.add_argument("--input-name", default="input", help="ONNX input name.") + parser.add_argument("--image-size", type=int, default=160, help="Square image size used by training.") + parser.add_argument("--samples", type=int, default=16, help="Number of calibration samples to write.") + return parser.parse_args() + + +def preprocess_image(path: Path, image_size: int) -> np.ndarray: + image = Image.open(path).convert("RGB").resize((image_size, image_size), Image.Resampling.BILINEAR) + array = np.asarray(image, dtype=np.float32) / 255.0 + array = np.transpose(array, (2, 0, 1)) + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None] + std = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None] + return ((array - mean) / std)[None, ...].astype("float32") + + +def main() -> None: + args = parse_args() + images = sorted(p for p in args.dataset_dir.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS) + if not images: + raise SystemExit(f"No images found under {args.dataset_dir}") + if args.samples < 1: + raise SystemExit("--samples must be at least 1") + + args.calibration_dir.mkdir(parents=True, exist_ok=True) + args.input_file.parent.mkdir(parents=True, exist_ok=True) + + sample_count = min(args.samples, len(images)) + prepared = [] + for index, image_path in enumerate(images[:sample_count]): + sample = preprocess_image(image_path, args.image_size) + np.save(args.calibration_dir / f"sample_{index:03d}.npy", sample) + prepared.append(sample) + + np.savez(args.input_file, **{args.input_name: prepared[0]}) + print(f"Wrote {sample_count} calibration samples to {args.calibration_dir}") + print(f"Wrote validation input to {args.input_file}") + + +if __name__ == "__main__": + main() diff --git a/examples/ai-hub/run_ai_hub.sh b/examples/ai-hub/run_ai_hub.sh new file mode 100755 index 0000000..07fcb49 --- /dev/null +++ b/examples/ai-hub/run_ai_hub.sh @@ -0,0 +1,156 @@ +#!/usr/bin/env bash +set -euo pipefail + +CONFIG_PATH="config.yaml" +CALIBRATION_PATH="examples/training/data/aihub_calibration" +INPUT_FILE="examples/training/data/inputs.npz" +FROM_STEP="quantize" +FROM_JOB="" +MODEL_S3_URI="" +ONNX_PATH="" +INPUT_NAME="" +DOWNLOAD=true +OUTPUT_PATH="" + +usage() { + cat <&2 + usage >&2 + exit 1 + ;; + esac +done + +if [[ ! -f "${CONFIG_PATH}" ]]; then + echo "Config not found: ${CONFIG_PATH}" >&2 + exit 1 +fi + +case "${FROM_STEP}" in + quantize|compile|validate|profile) + ;; + *) + echo "--from-step must be one of: quantize, compile, validate, profile" >&2 + exit 1 + ;; +esac + +if [[ ! -e "${CALIBRATION_PATH}" ]]; then + echo "Calibration path not found: ${CALIBRATION_PATH}" >&2 + echo "Pass --calibration with a .npz file or directory of .npy samples." >&2 + exit 1 +fi + +if [[ ! -f "${INPUT_FILE}" ]]; then + echo "Input file not found: ${INPUT_FILE}" >&2 + echo "Pass --input-file with a validation .npz or .npy file." >&2 + exit 1 +fi + +run() { + echo "+ $*" + "$@" +} + +UPLOAD_ARGS=( + "${CALIBRATION_PATH}" + "${INPUT_FILE}" + --from-step "${FROM_STEP}" + --config "${CONFIG_PATH}" +) + +if [[ -n "${FROM_JOB}" ]]; then + UPLOAD_ARGS+=(--from-job "${FROM_JOB}") +fi + +if [[ -n "${MODEL_S3_URI}" ]]; then + UPLOAD_ARGS+=(--model-s3-uri "${MODEL_S3_URI}") +fi + +if [[ -n "${ONNX_PATH}" ]]; then + UPLOAD_ARGS+=(--onnx-path "${ONNX_PATH}") +fi + +if [[ -n "${INPUT_NAME}" ]]; then + UPLOAD_ARGS+=(--input-name "${INPUT_NAME}") +fi + +run uv run qc-cli ai-hub upload "${UPLOAD_ARGS[@]}" + +if [[ "${DOWNLOAD}" == false ]]; then + exit 0 +fi + +DOWNLOAD_ARGS=(--config "${CONFIG_PATH}") +if [[ -n "${OUTPUT_PATH}" ]]; then + DOWNLOAD_ARGS+=(--output "${OUTPUT_PATH}") +fi + +run uv run qc-cli ai-hub download "${DOWNLOAD_ARGS[@]}" diff --git a/examples/training/source/train.py b/examples/training/source/train.py index 51c823e..6dd7a92 100644 --- a/examples/training/source/train.py +++ b/examples/training/source/train.py @@ -126,10 +126,6 @@ def export_onnx(model: nn.Module, model_dir: Path, image_size: int) -> None: do_constant_folding=True, input_names=["input"], output_names=["logits"], - dynamic_axes={ - "input": {0: "batch_size"}, - "logits": {0: "batch_size"}, - }, )