diff --git a/README.md b/README.md index 55e9bee..3a61e88 100644 --- a/README.md +++ b/README.md @@ -177,27 +177,35 @@ The expected output artifact is SageMaker’s `model.tar.gz`, normally containin ``` qc-cli ai-hub upload qc-cli ai-hub upload --from-step validate -qc-cli ai-hub quantize [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME] +qc-cli ai-hub optimize [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME] +qc-cli ai-hub quantize [--model-id ID] [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME] qc-cli ai-hub compile [--model-id ID] [--onnx-path PATH] [--model-s3-uri URI] [--from-job NAME] qc-cli ai-hub validate [--model-id ID] [--input-name NAME] qc-cli ai-hub profile [--model-id ID] qc-cli ai-hub download [--model-id ID] [--output PATH] ``` -`ai-hub upload` runs the four Workbench upload steps in order: quantize, compile, validate, and profile. Use `--from-step compile`, `--from-step validate`, or `--from-step profile` to resume from saved local state after a completed earlier step. +`ai-hub upload` optimizes to ONNX, quantizes, validates, and profiles. When `aihub.target_runtime` is not `onnx`, it +also compiles the quantized model to that deployment runtime. The initial ONNX optimization gives external models +Workbench provenance and applies compiler optimization passes before quantization. Resume behavior: ```text ---from-step quantize Run quantize, compile, validate, and profile. ---from-step compile Skip quantize; compile the last quantized model unless an explicit source is passed. ---from-step validate Skip quantize and compile; validate the last compiled model. ---from-step profile Skip quantize, compile, and validate; profile the last compiled model. +--from-step optimize Run optimize, quantize, optional final compile, validate, and profile. +--from-step quantize Quantize the last optimized ONNX, then optionally compile, validate, and profile. +--from-step compile Skip optimize and quantize; finalize the last quantized model for the target runtime. +--from-step validate Skip optimize, quantize, and compile; validate the last compiled model. +--from-step profile Skip optimize, quantize, compile, and validate; profile the last compiled model. ``` When a step runs in the current command, `upload` passes its returned model ID directly to the next step. When a step is skipped, the next step resolves the needed model ID from `.qc-cli.json`. This avoids re-running earlier AI Hub jobs when you only need to continue from a later step. -`ai-hub compile` resolves model sources in this order: `--model-id`, explicit source options (`--onnx-path`, `--model-s3-uri`, `--from-job`), last quantized model from state, then the last training job from local state. `ai-hub download` is separate because downloading the optimized artifact is outside the four-step Workbench upload loop. +`ai-hub optimize` compiles an external model with `--target_runtime onnx`. `ai-hub quantize` uses an explicit +`--model-id`, the last optimized ONNX model, or an explicit/local model source in that order. `ai-hub compile` resolves +model sources in this order: `--model-id`, explicit source options, last quantized model, then the last training job. +For `target_runtime: onnx`, upload treats the quantized ONNX as the final model and skips a redundant second compile. +`ai-hub download` remains separate because downloading is outside the Workbench processing loop. AI Hub authentication currently uses the local `qai-hub` SDK configuration. A planned follow-up is to support AWS Systems Manager Parameter Store `SecureString` for team-managed tokens, where `config.yaml` stores only a parameter name such as `/qc-cli/aihub/token`, AWS KMS encrypts the token at rest, and the CLI retrieves it at runtime with `ssm:GetParameter` plus `kms:Decrypt` permissions. diff --git a/examples/ai-hub/README.md b/examples/ai-hub/README.md deleted file mode 100644 index e0d865f..0000000 --- a/examples/ai-hub/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# 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 -examples/training/run_training.sh --wait -``` - -The `config.yaml` file must include AI Hub settings: - -```yaml -aihub: - device: - name: Samsung Galaxy S25 (Family) - target_runtime: tflite - input_specs: - input: [[1, 3, 160, 160], float32] - output_dir: build/qai-hub -``` - -Finally, the user needs to authenticate with Qualcomm AI Hub using: - -```bash -qai-hub configure --api_token -``` - -## 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. - -To generate calibration and validation inputs: - -```bash -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 - -## Upload Model to Qualcomm Workbench - -The model can be uploaded to Qualcomm Workbench using: - -```bash -qc-cli ai-hub upload examples/training/data/aihub_calibration examples/training/data/inputs.npz -``` - -The first argument is the calibration path for the model and the second argument is the input file, both of which were created by the `prepare_inputs.py` script. For more details, add `--help` after the `upload` command. - -The `upload` command runs the following commands in order: -1. `qc-cli ai-hub quantize` -2. `qc-cli ai-hub compile` -3. `qc-cli ai-hub validate` -4. `qc-cli ai-hub profile` - -Finally the user can download the model from AI Workbench using the command -```bash -qc-cli ai-hub download -``` diff --git a/examples/ai-hub/prepare_inputs.py b/examples/ai-hub/prepare_inputs.py deleted file mode 100644 index 72514c7..0000000 --- a/examples/ai-hub/prepare_inputs.py +++ /dev/null @@ -1,74 +0,0 @@ -#!/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/meter-detection/README.md b/examples/meter-detection/README.md new file mode 100644 index 0000000..35955df --- /dev/null +++ b/examples/meter-detection/README.md @@ -0,0 +1,266 @@ +# YOLO26 Electric Meter Detection Example + +This example trains a YOLO26 object detection model on the Roboflow Universe electric meter dataset using the existing `qc-cli` SageMaker training flow. + +The workflow is intentionally command driven. Run each step yourself so you can inspect the dataset, update `config.yaml`, and decide when to submit the SageMaker job. + +Dataset: + +```text +https://universe.roboflow.com/kemals-workspace-kbc8l/electric-meter-detection-o4tfi/dataset/1 +``` + +## Prerequisites + +- Install or sync the project dependencies: `uv sync` +- The virtual environment is activated. +- AWS credentials configured for the profile in `config.yaml` +- Infrastructure already deployed with `qc-cli infra setup` + +## 1. Download The Dataset + +Register or sign in to Roboflow, then open the dataset page: + +```text +https://universe.roboflow.com/kemals-workspace-kbc8l/electric-meter-detection-o4tfi/dataset/1 +``` + +Download the dataset in YOLOv26 format from the Roboflow UI, then extract the downloaded archive into: + +```text +examples/meter-detection/data/electric-meter-detection +``` + +The `data.yaml` file should be directly under that folder: + +```text +examples/meter-detection/data/electric-meter-detection/data.yaml +``` + +Do not move `data.yaml` into the `train/` split folder. + +After extracting, confirm the dataset has a YOLO data file and image splits: + +```bash +find examples/meter-detection/data/electric-meter-detection -maxdepth 2 -type d | sort +find examples/meter-detection/data/electric-meter-detection -name data.yaml -print +``` + +Open `examples/meter-detection/data/electric-meter-detection/data.yaml` and make sure the split paths are relative to that folder: + +```yaml +path: . +train: train/images +val: valid/images +test: test/images +``` + +If your downloaded dataset does not include a `test/` folder, remove the `test:` line. + +The expected layout is similar to: + +```text +examples/meter-detection/data/electric-meter-detection/ + data.yaml + train/ + valid/ + test/ +``` + +## 2. Configure SageMaker Training + +Update `config.yaml` so the training section points at this example's source directory: + +```yaml +sagemaker: + training: + image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1 + instance_type: ml.g4dn.xlarge + instance_count: 1 + source_dir: examples/meter-detection/source + entry_point: train.py + hyperparameters: + model: yolo26n.pt + epochs: 25 + imgsz: 640 + batch: 16 + workers: 2 +``` + +Use `yolo26n.pt` for a lightweight first YOLO26 run. If those weights are unavailable in the installed Ultralytics package, use `yolo11n.pt` as the established fallback: + +```yaml + model: yolo11n.pt +``` + +The `source/requirements.txt` file is installed by the SageMaker PyTorch container before running `train.py`. + +For a CPU smoke test, use a CPU instance and reduce the workload: + +```yaml +sagemaker: + training: + image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1 + instance_type: ml.m4.xlarge + instance_count: 1 + source_dir: examples/meter-detection/source + entry_point: train.py + hyperparameters: + model: yolo26n.pt + epochs: 1 + imgsz: 320 + batch: 4 + workers: 2 +``` + +## 3. Check Infrastructure + +Confirm the CLI can see the configured SageMaker role and S3 bucket: + +```bash +qc-cli infra status +``` + +## 4. Upload The Dataset + +Upload the downloaded Roboflow dataset to the `s3.data_prefix` configured in `config.yaml`: + +```bash +qc-cli upload examples/meter-detection/data/electric-meter-detection +``` + +Directory uploads preserve paths relative to the uploaded directory, so SageMaker receives the dataset root with `data.yaml` plus the split directories. + +In SageMaker, this uploaded dataset root is mounted at `/opt/ml/input/data/train`. That `train` path is the SageMaker channel name, not the YOLO `train/` split folder. + +## 5. Start Training + +Submit the SageMaker training job: + +```bash +qc-cli train start +``` + +The command prints the submitted SageMaker job name. Check progress with: + +```bash +qc-cli train status +``` + +Or pass the job name explicitly: + +```bash +qc-cli train status qc-cli-YYYYMMDD-HHMMSS +``` + +## SageMaker Outputs + +When the job completes, SageMaker packages the files written under `/opt/ml/model` into `model.tar.gz`. + +This example writes: + +```text +best.pt +model.onnx +metrics.json +``` + +The archive is stored under the configured `s3.model_prefix`. + +## 6. Configure Qualcomm AI Hub + +Authenticate with Qualcomm AI Hub: + +```bash +qai-hub configure --api_token +``` + +Add AI Hub settings to `config.yaml`. The input name and image size must match the ONNX model exported by this example: + +```yaml +aihub: + device: + name: Dragonwing IQ-9075 EVK + target_runtime: onnx + input_specs: + images: [[1, 3, 640, 640], float32] + job_name: meter-detection + model_name: meter-detection + output_dir: build/qai-hub/meter-detection +``` + +The ONNX graph is the source of truth. The export normally uses the same value as `sagemaker.training.hyperparameters.imgsz`, but changing `config.yaml` after training does not resize an existing model. For example, a model exported with `imgsz: 320` requires `images: [[1, 3, 320, 320], float32]`. + +## 7. Prepare AI Hub Inputs + +Generate calibration samples and a validation input from the downloaded dataset: + +```bash +uv run python examples/meter-detection/prepare_aihub_inputs.py --image-size 640 +``` + +This writes: + +```text +examples/meter-detection/data/aihub_calibration/*.npy +examples/meter-detection/data/inputs.npz +``` + +The script applies the preprocessing expected by the exported YOLO model: aspect-ratio-preserving letterboxing, RGB channel order, channel-first layout, and pixel values normalized to `[0, 1]`. + +## 8. Upload To Qualcomm AI Hub + +Use the SageMaker job name printed by `qc-cli train start`: + +```bash +qc-cli ai-hub upload \ + examples/meter-detection/data/aihub_calibration \ + examples/meter-detection/data/inputs.npz \ + --from-job qc-cli-YYYYMMDD-HHMMSS +``` + +The command downloads the job's `model.tar.gz`, finds `model.onnx`, and runs the following AI Hub workflow: + +1. Compile the external ONNX to a Workbench-optimized ONNX model. +2. Quantize the optimized ONNX model. +3. Compile the quantized model when the configured deployment runtime is not `onnx`. +4. Validate and profile the final model. + +The training example sanitizes the Ultralytics ONNX export before saving `model.onnx`. This removes graph input or output names, such as `output0`, that are duplicated in the ONNX `value_info` metadata and rejected by AI Hub. + +For a model already downloaded by a failed upload attempt, sanitize the extracted ONNX file and retry using the local model. Replace the job name in both paths: + +```bash +uv run --with onnx python examples/meter-detection/source/sanitize_onnx.py \ + build/qai-hub/meter-detection/qc-cli-YYYYMMDD-HHMMSS/source/extracted/model.onnx \ + --output build/qai-hub/meter-detection/model.aihub.onnx + +qc-cli ai-hub upload \ + examples/meter-detection/data/aihub_calibration \ + examples/meter-detection/data/inputs.npz \ + --onnx-path build/qai-hub/meter-detection/model.aihub.onnx +``` + +Download the compiled artifact after the workflow completes: + +```bash +qc-cli ai-hub download --output build/qai-hub/meter-detection/model.tflite +``` + +## Training Hyperparameters + +Values under `sagemaker.training.hyperparameters` are passed to `source/train.py` as command-line arguments. + +| Name | Type | Default | Description | +|---|---:|---:|---| +| `model` | string | `yolo26n.pt` | Ultralytics model weights or model YAML. | +| `epochs` | int | `25` | Number of training epochs. | +| `imgsz` | int | `640` | Square training image size. | +| `batch` | int | `16` | Images per training batch. | +| `workers` | int | `2` | DataLoader worker count. | +| `patience` | int | `20` | Early stopping patience. | +| `device` | string | auto | Optional Ultralytics device value such as `0` or `cpu`. | +| `data-yaml` | string | auto | Optional path to `data.yaml`; normally discovered from the uploaded dataset root. | +| `dataset-dir` | string | `SM_CHANNEL_TRAIN` | Uploaded dataset root mounted by SageMaker. | + +Do not set `dataset-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`. diff --git a/examples/meter-detection/prepare_aihub_inputs.py b/examples/meter-detection/prepare_aihub_inputs.py new file mode 100644 index 0000000..6b0e5a5 --- /dev/null +++ b/examples/meter-detection/prepare_aihub_inputs.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Prepare Qualcomm AI Hub calibration and validation inputs for the meter detector.""" + +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/meter-detection/data/electric-meter-detection"), + help="Root of the extracted Roboflow dataset.", + ) + parser.add_argument( + "--calibration-dir", + type=Path, + default=Path("examples/meter-detection/data/aihub_calibration"), + help="Directory where .npy calibration samples will be written.", + ) + parser.add_argument( + "--input-file", + type=Path, + default=Path("examples/meter-detection/data/inputs.npz"), + help="Validation .npz input file for qc-cli ai-hub validate.", + ) + parser.add_argument("--input-name", default="images", help="ONNX input name.") + parser.add_argument("--image-size", type=int, default=640, help="Square image size used for ONNX export.") + 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: + """Apply Ultralytics-style letterboxing and produce an NCHW float32 tensor.""" + with Image.open(path) as source: + image = source.convert("RGB") + + scale = min(image_size / image.width, image_size / image.height) + resized_width = round(image.width * scale) + resized_height = round(image.height * scale) + image = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR) + + canvas = Image.new("RGB", (image_size, image_size), (114, 114, 114)) + left = round((image_size - resized_width) / 2 - 0.1) + top = round((image_size - resized_height) / 2 - 0.1) + canvas.paste(image, (left, top)) + + array = np.asarray(canvas, dtype=np.float32) / 255.0 + return np.transpose(array, (2, 0, 1))[None, ...].astype(np.float32) + + +def main() -> None: + args = parse_args() + if args.image_size < 1: + raise SystemExit("--image-size must be at least 1") + if args.samples < 1: + raise SystemExit("--samples must be at least 1") + + images = sorted( + path + for path in args.dataset_dir.rglob("*") + if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS and path.parent.name == "images" + ) + if not images: + raise SystemExit(f"No images found under {args.dataset_dir}") + + args.calibration_dir.mkdir(parents=True, exist_ok=True) + args.input_file.parent.mkdir(parents=True, exist_ok=True) + for stale_sample in args.calibration_dir.glob("sample_*.npy"): + stale_sample.unlink() + + prepared: list[np.ndarray] = [] + for index, image_path in enumerate(images[: args.samples]): + 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]}) # pyright: ignore[reportArgumentType] + print(f"Wrote {len(prepared)} calibration samples to {args.calibration_dir}") + print(f"Wrote validation input to {args.input_file}") + + +if __name__ == "__main__": + main() diff --git a/examples/meter-detection/source/requirements.txt b/examples/meter-detection/source/requirements.txt new file mode 100644 index 0000000..419305a --- /dev/null +++ b/examples/meter-detection/source/requirements.txt @@ -0,0 +1,3 @@ +ultralytics>=8.3.0 +pyyaml>=6.0.3 +onnx>=1.16.0 diff --git a/examples/meter-detection/source/sanitize_onnx.py b/examples/meter-detection/source/sanitize_onnx.py new file mode 100644 index 0000000..f7d1103 --- /dev/null +++ b/examples/meter-detection/source/sanitize_onnx.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +"""Remove ONNX value_info entries that duplicate graph inputs or outputs.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import onnx # type: ignore[reportMissingImports] + + +def sanitize_onnx(path: Path, output_path: Path | None = None) -> Path: + model = onnx.load(path) + io_names = {value.name for value in (*model.graph.input, *model.graph.output)} + retained_value_info = [value for value in model.graph.value_info if value.name not in io_names] + + destination = output_path or path + if len(retained_value_info) != len(model.graph.value_info): + del model.graph.value_info[:] + model.graph.value_info.extend(retained_value_info) + + destination.parent.mkdir(parents=True, exist_ok=True) + onnx.save(model, destination) + return destination + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("onnx_path", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + written = sanitize_onnx(args.onnx_path, args.output) + print(f"Saved sanitized ONNX model to {written}") + + +if __name__ == "__main__": + main() diff --git a/examples/meter-detection/source/train.py b/examples/meter-detection/source/train.py new file mode 100644 index 0000000..99aef51 --- /dev/null +++ b/examples/meter-detection/source/train.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""SageMaker entry point for YOLO electric meter detection training.""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from pathlib import Path +from typing import Any + +import yaml +from sanitize_onnx import sanitize_onnx +from ultralytics import YOLO # type: ignore[reportMissingImports] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model", default="yolo26n.pt") + parser.add_argument("--epochs", type=int, default=25) + parser.add_argument("--imgsz", type=int, default=640) + parser.add_argument("--batch", type=int, default=16) + parser.add_argument("--workers", type=int, default=2) + parser.add_argument("--patience", type=int, default=20) + parser.add_argument("--device", default=None) + parser.add_argument("--data-yaml", default=None) + parser.add_argument("--dataset-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train")) + parser.add_argument("--train-dir", dest="dataset_dir", help=argparse.SUPPRESS) + parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")) + return parser.parse_args() + + +def find_data_yaml(dataset_dir: Path, explicit_path: str | None) -> Path: + if explicit_path: + data_yaml = Path(explicit_path) + if data_yaml.is_file(): + return data_yaml + raise FileNotFoundError(f"Configured data.yaml does not exist: {data_yaml}") + + matches = sorted(dataset_dir.rglob("data.yaml")) + if not matches: + raise FileNotFoundError(f"Could not find data.yaml under {dataset_dir}") + if len(matches) > 1: + print(f"Found multiple data.yaml files; using {matches[0]}") + return matches[0] + + +def prepare_data_yaml(data_yaml: Path) -> Path: + """Write a SageMaker-local data file rooted at the uploaded dataset.""" + dataset_root = data_yaml.parent + data = yaml.safe_load(data_yaml.read_text(encoding="utf-8")) + if not isinstance(data, dict): + raise ValueError(f"Expected a mapping in {data_yaml}") + + normalized = dict(data) + normalized["path"] = str(dataset_root) + if "val" not in normalized and "valid" in normalized: + normalized["val"] = normalized.pop("valid") + + prepared_path = dataset_root / "data.sagemaker.yaml" + prepared_path.write_text(yaml.safe_dump(normalized, sort_keys=False), encoding="utf-8") + print(f"Prepared dataset config: {prepared_path}") + return prepared_path + + +def copy_if_exists(source: Path, destination: Path) -> None: + if source.exists(): + shutil.copy2(source, destination) + print(f"Saved {destination}") + + +def main() -> None: + args = parse_args() + dataset_dir = Path(args.dataset_dir) + model_dir = Path(args.model_dir) + model_dir.mkdir(parents=True, exist_ok=True) + + data_yaml = prepare_data_yaml(find_data_yaml(dataset_dir, args.data_yaml)) + model = YOLO(args.model) + + train_kwargs: dict[str, Any] = { + "data": str(data_yaml), + "epochs": args.epochs, + "imgsz": args.imgsz, + "batch": args.batch, + "workers": args.workers, + "patience": args.patience, + "project": str(model_dir / "runs"), + "name": "train", + "exist_ok": True, + } + if args.device: + train_kwargs["device"] = args.device + + results = model.train(**train_kwargs) + save_dir = Path(results.save_dir) + best_pt = save_dir / "weights" / "best.pt" + last_pt = save_dir / "weights" / "last.pt" + trained_weights = best_pt if best_pt.exists() else last_pt + if not trained_weights.exists(): + raise FileNotFoundError(f"Could not find trained weights in {save_dir / 'weights'}") + + copy_if_exists(trained_weights, model_dir / "best.pt") + trained_model = YOLO(str(trained_weights)) + onnx_path = Path(trained_model.export(format="onnx", imgsz=args.imgsz)) + saved_onnx_path = sanitize_onnx(onnx_path, model_dir / "model.onnx") + print(f"Saved {saved_onnx_path}") + + metrics = { + "model": args.model, + "epochs": args.epochs, + "imgsz": args.imgsz, + "batch": args.batch, + "workers": args.workers, + "patience": args.patience, + "data_yaml": str(data_yaml), + "weights": str(trained_weights), + "onnx": str(saved_onnx_path), + } + (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") + print(f"Saved model artifacts to {model_dir}") + + +if __name__ == "__main__": + main() diff --git a/examples/training/README.md b/examples/training/README.md deleted file mode 100644 index d968609..0000000 --- a/examples/training/README.md +++ /dev/null @@ -1,89 +0,0 @@ -# SageMaker Training Example - -This example downloads a small image-classification dataset, uploads it through `qc-cli`, and submits a live SageMaker training job. - -## Prerequisites - -- AWS credentials configured for the profile in `config.yaml` -- Infrastructure already deployed with `qc-cli infra setup` -- `config.yaml` updated with: - -```yaml -s3: - bucket: your-bucket-name - -sagemaker: - training: - image_uri: 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.6-cpu-py312-ubuntu22.04-sagemaker-v1 - instance_type: ml.m4.xlarge - instance_count: 1 - source_dir: examples/training/source - entry_point: train.py - hyperparameters: - epochs: 1 - batch-size: 32 - learning-rate: 0.001 - image-size: 160 - validation-split: 0.2 -``` - -## Training Hyperparameters - -Values under `sagemaker.training.hyperparameters` are passed to the training entry point as command-line arguments. For this example, they map to arguments defined in [source/train.py](source/train.py). - -Supported by this example: - -| Name | Type | Default | Description | -|---|---:|---:|---| -| `epochs` | int | `1` | Number of training epochs. | -| `batch-size` | int | `32` | Images per training batch. | -| `learning-rate` | float | `0.001` | Adam optimizer learning rate. | -| `image-size` | int | `160` | Resize images to square `image-size x image-size`. | -| `validation-split` | float | `0.2` | Fraction of data used for validation. | -| `max-samples` | int | `0` | Optional cap for smoke tests; `0` means use all images. | -| `seed` | int | `13` | Random seed for reproducible splitting. | -| `num-workers` | int | `2` | DataLoader worker count. | - -Do not set `train-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`. - -## 1. Download The Dataset - -```bash -bash examples/training/download_flower_photos.sh -``` - -This creates: - -```text -examples/training/data/flower_photos_sagemaker/ - daisy/ - dandelion/ - roses/ - sunflowers/ - tulips/ -``` - -## 2. Run Training - -Run the training script and wait until it finishes: - -```bash -bash examples/training/run_training.sh --config config.yaml --wait -``` - -Use a dataset that is already uploaded to `s3.data_prefix`: - -```bash -bash examples/training/run_training.sh \ - --config config.yaml \ - --skip-upload \ - --wait -``` - -## Notes - -- The default dataset path is `examples/training/data/flower_photos_sagemaker`. -- Uploaded data uses the `s3.bucket` and `s3.data_prefix` values from `config.yaml`. -- Training artifacts are written under `s3:////`. -- The SageMaker `model.tar.gz` contains `model.onnx`, `model.pt`, `class_to_idx.json`, and `metrics.json`. -- SageMaker packages `examples/training/source`, installs `requirements.txt`, and runs `train.py`. diff --git a/examples/training/download_flower_photos.sh b/examples/training/download_flower_photos.sh deleted file mode 100755 index 5df0429..0000000 --- a/examples/training/download_flower_photos.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -DATASET_URL="https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz" -DEST_DIR="${1:-examples/training/data}" -ARCHIVE_PATH="${DEST_DIR}/flower_photos.tgz" -RAW_DATASET_DIR="${DEST_DIR}/flower_photos" -DATASET_DIR="${DEST_DIR}/flower_photos_sagemaker" -CLASS_NAMES=("daisy" "dandelion" "roses" "sunflowers" "tulips") - -mkdir -p "${DEST_DIR}" - -if [[ -d "${DATASET_DIR}" ]]; then - echo "Dataset already exists: ${DATASET_DIR}" - echo "Use this path with run_training.py:" - echo " ${DATASET_DIR}" - exit 0 -fi - -echo "Downloading TensorFlow flower_photos dataset..." -if command -v curl >/dev/null 2>&1; then - curl -L "${DATASET_URL}" -o "${ARCHIVE_PATH}" -elif command -v wget >/dev/null 2>&1; then - wget -O "${ARCHIVE_PATH}" "${DATASET_URL}" -else - echo "Either curl or wget is required." >&2 - exit 1 -fi - -echo "Extracting dataset..." -tar -xzf "${ARCHIVE_PATH}" -C "${DEST_DIR}" - -echo "Preparing SageMaker directory layout..." -mkdir -p "${DATASET_DIR}" -for class_name in "${CLASS_NAMES[@]}"; do - cp -R "${RAW_DATASET_DIR}/${class_name}" "${DATASET_DIR}/${class_name}" -done - -echo "Dataset ready: ${DATASET_DIR}" -find "${DATASET_DIR}" -mindepth 1 -maxdepth 1 -type d -print | sort diff --git a/examples/training/run_training.sh b/examples/training/run_training.sh deleted file mode 100755 index 55eb3bb..0000000 --- a/examples/training/run_training.sh +++ /dev/null @@ -1,112 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -CONFIG_PATH="config.yaml" -DATASET_DIR="examples/training/data/flower_photos_sagemaker" -WAIT=false -SKIP_UPLOAD=false -POLL_SECONDS=60 - -usage() { - cat <&2 - usage >&2 - exit 1 - ;; - esac -done - -if [[ ! -f "${CONFIG_PATH}" ]]; then - echo "Config not found: ${CONFIG_PATH}" >&2 - exit 1 -fi - -if [[ "${SKIP_UPLOAD}" == false && ! -d "${DATASET_DIR}" ]]; then - echo "Dataset not found: ${DATASET_DIR}" >&2 - echo "Run: bash examples/training/download_flower_photos.sh" >&2 - exit 1 -fi - -run() { - echo "+ $*" - "$@" -} - -run uv run qc-cli infra status --config "${CONFIG_PATH}" - -if [[ "${SKIP_UPLOAD}" == false ]]; then - run uv run qc-cli upload "${DATASET_DIR}" --config "${CONFIG_PATH}" -fi - -TRAIN_OUTPUT_FILE="$(mktemp)" -trap 'rm -f "${TRAIN_OUTPUT_FILE}"' EXIT -run uv run qc-cli train start --config "${CONFIG_PATH}" | tee "${TRAIN_OUTPUT_FILE}" - -JOB_NAME="$(grep -Eo 'qc-cli-[0-9]{8}-[0-9]{6}' "${TRAIN_OUTPUT_FILE}" | tail -n 1)" -if [[ -z "${JOB_NAME}" ]]; then - echo "Could not find training job name in qc-cli output." >&2 - exit 1 -fi - -echo "Submitted SageMaker training job: ${JOB_NAME}" - -if [[ "${WAIT}" == false ]]; then - run uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}" - exit 0 -fi - -while true; do - STATUS_OUTPUT="$(uv run qc-cli train status "${JOB_NAME}" --config "${CONFIG_PATH}")" - echo "${STATUS_OUTPUT}" - - if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Completed'; then - echo "Training completed successfully." - exit 0 - fi - - if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Failed'; then - echo "Training failed." >&2 - exit 1 - fi - - if printf '%s\n' "${STATUS_OUTPUT}" | grep -q 'Status:.*Stopped'; then - echo "Training stopped." >&2 - exit 1 - fi - - sleep "${POLL_SECONDS}" -done diff --git a/examples/training/source/requirements.txt b/examples/training/source/requirements.txt deleted file mode 100644 index 90dcba3..0000000 --- a/examples/training/source/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -onnx==1.21.0 diff --git a/examples/training/source/train.py b/examples/training/source/train.py deleted file mode 100644 index 6dd7a92..0000000 --- a/examples/training/source/train.py +++ /dev/null @@ -1,188 +0,0 @@ -#!/usr/bin/env python3 -"""SageMaker entry point for CPU image-classification training.""" - -from __future__ import annotations - -import argparse -import json -import os -import random -from pathlib import Path - -import torch -from torch import nn -from torch.utils.data import DataLoader, Subset, random_split -from torchvision import datasets, transforms - - -class SmallImageClassifier(nn.Module): - def __init__(self, class_count: int) -> None: - super().__init__() - self.features = nn.Sequential( - nn.Conv2d(3, 16, kernel_size=3, padding=1), - nn.ReLU(inplace=True), - nn.MaxPool2d(2), - nn.Conv2d(16, 32, kernel_size=3, padding=1), - nn.ReLU(inplace=True), - nn.MaxPool2d(2), - nn.Conv2d(32, 64, kernel_size=3, padding=1), - nn.ReLU(inplace=True), - nn.MaxPool2d(2), - nn.AdaptiveAvgPool2d((1, 1)), - ) - self.classifier = nn.Linear(64, class_count) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.features(x) - x = torch.flatten(x, 1) - return self.classifier(x) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--epochs", type=int, default=1) - parser.add_argument("--batch-size", type=int, default=32) - parser.add_argument("--learning-rate", type=float, default=0.001) - parser.add_argument("--image-size", type=int, default=160) - parser.add_argument("--validation-split", type=float, default=0.2) - parser.add_argument("--max-samples", type=int, default=0) - parser.add_argument("--seed", type=int, default=13) - parser.add_argument("--num-workers", type=int, default=2) - parser.add_argument("--train-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train")) - parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model")) - return parser.parse_args() - - -def build_datasets(args: argparse.Namespace) -> tuple[Subset, Subset, dict[str, int]]: - transform = transforms.Compose( - [ - transforms.Resize((args.image_size, args.image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), - ] - ) - dataset = datasets.ImageFolder(args.train_dir, transform=transform) - if len(dataset.classes) < 2: - raise ValueError(f"Expected at least two classes in {args.train_dir}. Found: {dataset.classes}") - - if args.max_samples > 0 and args.max_samples < len(dataset): - indices = list(range(len(dataset))) - random.Random(args.seed).shuffle(indices) - dataset = Subset(dataset, indices[: args.max_samples]) - - validation_size = max(1, int(len(dataset) * args.validation_split)) - train_size = len(dataset) - validation_size - if train_size < 1: - raise ValueError("Not enough images to create a train/validation split.") - - generator = torch.Generator().manual_seed(args.seed) - train_dataset, validation_dataset = random_split(dataset, [train_size, validation_size], generator=generator) - return train_dataset, validation_dataset, getattr(dataset, "dataset", dataset).class_to_idx - - -def run_epoch( - model: nn.Module, - data_loader: DataLoader, - criterion: nn.Module, - optimizer: torch.optim.Optimizer | None, - device: torch.device, -) -> tuple[float, float]: - training = optimizer is not None - model.train(training) - - total_loss = 0.0 - total_correct = 0 - total_examples = 0 - - for images, labels in data_loader: - images = images.to(device) - labels = labels.to(device) - - with torch.set_grad_enabled(training): - logits = model(images) - loss = criterion(logits, labels) - - if training: - optimizer.zero_grad() - loss.backward() - optimizer.step() - - total_loss += loss.item() * images.size(0) - total_correct += (logits.argmax(dim=1) == labels).sum().item() - total_examples += images.size(0) - - return total_loss / total_examples, total_correct / total_examples - - -def export_onnx(model: nn.Module, model_dir: Path, image_size: int) -> None: - model.eval() - dummy_input = torch.randn(1, 3, image_size, image_size) - torch.onnx.export( - model, - dummy_input, - model_dir / "model.onnx", - export_params=True, - opset_version=17, - do_constant_folding=True, - input_names=["input"], - output_names=["logits"], - ) - - -def main() -> None: - args = parse_args() - random.seed(args.seed) - torch.manual_seed(args.seed) - - train_dataset, validation_dataset, class_to_idx = build_datasets(args) - train_loader = DataLoader( - train_dataset, - batch_size=args.batch_size, - shuffle=True, - num_workers=args.num_workers, - ) - validation_loader = DataLoader( - validation_dataset, - batch_size=args.batch_size, - shuffle=False, - num_workers=args.num_workers, - ) - - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - model = SmallImageClassifier(class_count=len(class_to_idx)).to(device) - criterion = nn.CrossEntropyLoss() - optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) - - print(f"Training on {device}. Classes: {sorted(class_to_idx)}") - metrics = [] - for epoch in range(1, args.epochs + 1): - train_loss, train_accuracy = run_epoch(model, train_loader, criterion, optimizer, device) - validation_loss, validation_accuracy = run_epoch(model, validation_loader, criterion, None, device) - epoch_metrics = { - "epoch": epoch, - "train_loss": train_loss, - "train_accuracy": train_accuracy, - "validation_loss": validation_loss, - "validation_accuracy": validation_accuracy, - } - metrics.append(epoch_metrics) - print(json.dumps(epoch_metrics, sort_keys=True)) - - model_dir = Path(args.model_dir) - model_dir.mkdir(parents=True, exist_ok=True) - torch.save( - { - "model_state_dict": model.cpu().state_dict(), - "class_to_idx": class_to_idx, - "image_size": args.image_size, - }, - model_dir / "model.pt", - ) - export_onnx(model, model_dir, args.image_size) - (model_dir / "class_to_idx.json").write_text(json.dumps(class_to_idx, indent=2), encoding="utf-8") - (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") - print(f"Saved model artifacts to {model_dir}") - - -if __name__ == "__main__": - main() diff --git a/src/commands/ai_hub.py b/src/commands/ai_hub.py index ec80d58..1d0aeeb 100644 --- a/src/commands/ai_hub.py +++ b/src/commands/ai_hub.py @@ -1,4 +1,5 @@ from collections.abc import Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from enum import StrEnum from pathlib import Path @@ -12,9 +13,9 @@ from src import state as state_ops from src.commands.utils import CONFIG_OPT, CONSOLE, load_cfg from src.config import Config from src.qualcomm import aihub_jobs -from src.qualcomm.artifacts import resolve_onnx +from src.qualcomm.artifacts import ResolvedOnnx, resolve_onnx -app = typer.Typer(help="Quantize, compile, validate, profile, and download models with Qualcomm Workbench") +app = typer.Typer(help="Optimize, quantize, compile, validate, profile, and download models with Qualcomm Workbench") _RUNTIME_EXTENSIONS = { "tflite": "tflite", @@ -24,12 +25,19 @@ _RUNTIME_EXTENSIONS = { class UploadStep(StrEnum): + optimize = "optimize" quantize = "quantize" compile = "compile" validate = "validate" profile = "profile" +@dataclass(frozen=True) +class ResolvedModelSource: + model: str | Path + model_artifact: str | None = None + + def _input_specs(cfg: Config) -> dict[str, tuple[tuple[int, ...], str]]: specs = {name: (tuple(shape), dtype) for name, (shape, dtype) in cfg.aihub.input_specs.items()} if not specs: @@ -101,6 +109,57 @@ def _model_id_or_state(config_path: str, model_id: str | None, *, quantized: boo return resolved +def _resolve_model_source( + cfg: Config, + config_path: str, + *, + model_id: str | None = None, + previous_model_id: str | None = None, + from_job: str | None = None, + model_s3_uri: str | None = None, + onnx_path: str | None = None, +) -> ResolvedModelSource: + if model_id: + return ResolvedModelSource(model_id) + + has_explicit_source = bool(from_job or model_s3_uri or onnx_path) + if previous_model_id and not has_explicit_source: + return ResolvedModelSource(previous_model_id) + + resolved = _resolve_onnx_source( + cfg, + config_path, + from_job=from_job, + model_s3_uri=model_s3_uri, + onnx_path=onnx_path, + ) + return ResolvedModelSource(resolved.onnx_path, resolved.model_artifact) + + +def _resolve_onnx_source( + cfg: Config, + config_path: str, + *, + from_job: str | None = None, + model_s3_uri: str | None = None, + onnx_path: str | None = None, +) -> ResolvedOnnx: + st = state_ops.store(config_path) + last_training_job = st.get_last_training_job() + saved_model_artifact = None + if not from_job and not model_s3_uri and not onnx_path and not last_training_job: + saved_model_artifact = st.get_last_model_artifact() + + return resolve_onnx( + cfg=cfg, + output_dir=cfg.aihub.output_dir, + from_job=from_job, + model_s3_uri=model_s3_uri or saved_model_artifact, + onnx_path=onnx_path, + last_training_job=last_training_job, + ) + + def _device_selector(device: Device) -> str: parts: list[str] = [] if device.name: @@ -132,20 +191,23 @@ def _quantize_step( cfg: Config, config_path: str, calibration_path: Path, - from_job: str | None, - model_s3_uri: str | None, - onnx_path: str | None, + *, + model_id: str | None = None, + from_job: str | None = None, + model_s3_uri: str | None = None, + onnx_path: str | None = None, ) -> str: st = state_ops.store(config_path) specs = _input_specs(cfg) try: - resolved = resolve_onnx( - cfg=cfg, - output_dir=cfg.aihub.output_dir, + source = _resolve_model_source( + cfg, + config_path, + model_id=model_id, + previous_model_id=st.get_last_optimized_model_id(), from_job=from_job, - model_s3_uri=model_s3_uri or st.get_last_model_artifact(), + model_s3_uri=model_s3_uri, onnx_path=onnx_path, - last_training_job=st.get_last_training_job(), ) calibration_data = _load_calibration(calibration_path, specs) except (FileNotFoundError, ValueError) as e: @@ -153,73 +215,117 @@ def _quantize_step( raise typer.Exit(1) try: + hub_model = ( + hub.upload_model(str(source.model), name=cfg.aihub.model_name) + if isinstance(source.model, Path) + else hub.get_model(source.model) + ) result = aihub_jobs.submit_quantize_job( - resolved.onnx_path, + hub_model, calibration_data, cfg.aihub.quantize_options, job_name=_job_name(cfg, "quantize"), - model_name=cfg.aihub.model_name, ) except Exception as e: CONSOLE.print(f"[red]AI Hub quantize failed: {e}[/red]") raise typer.Exit(1) - st.update( - last_model_artifact=resolved.model_artifact, - last_quantize_job_id=result["job_id"], - last_quantized_model_id=result["model_id"], - ) + updates: dict[str, Any] = { + "last_quantize_job_id": result["job_id"], + "last_quantized_model_id": result["model_id"], + } + if source.model_artifact: + updates["last_model_artifact"] = source.model_artifact + st.update(**updates) CONSOLE.print(f"[green]✓[/green] Quantize job: [bold]{result['job_id']}[/bold]") CONSOLE.print(f"[green]✓[/green] Quantized model: [bold]{result['model_id']}[/bold]") return str(result["model_id"]) +def _optimize_step( + cfg: Config, + config_path: str, + from_job: str | None, + model_s3_uri: str | None, + onnx_path: str | None, +) -> str: + st = state_ops.store(config_path) + _validate_device(cfg) + specs = _input_specs(cfg) + try: + source = _resolve_onnx_source( + cfg, + config_path, + from_job=from_job, + model_s3_uri=model_s3_uri, + onnx_path=onnx_path, + ) + except (FileNotFoundError, ValueError) as e: + CONSOLE.print(f"[red]{e}[/red]") + raise typer.Exit(1) + + try: + hub_model = hub.upload_model(str(source.onnx_path), name=cfg.aihub.model_name) + result = aihub_jobs.submit_compile_job( + model=hub_model, + device=cfg.aihub.device, + input_specs=specs, + target_runtime="onnx", + job_name=_job_name(cfg, "optimize"), + ) + except Exception as e: + CONSOLE.print(f"[red]AI Hub ONNX optimization failed: {e}[/red]") + raise typer.Exit(1) + + st.update( + last_model_artifact=source.model_artifact, + last_optimize_job_id=result["job_id"], + last_optimized_model_id=result["model_id"], + ) + CONSOLE.print(f"[green]✓[/green] ONNX optimization job: [bold]{result['job_id']}[/bold]") + CONSOLE.print(f"[green]✓[/green] Optimized ONNX model: [bold]{result['model_id']}[/bold]") + return str(result["model_id"]) + + def _compile_step( cfg: Config, config_path: str, - model_id: str | None, - from_job: str | None, - model_s3_uri: str | None, - onnx_path: str | None, *, - prefer_quantized: bool, + model_id: str | None = None, + from_job: str | None = None, + model_s3_uri: str | None = None, + onnx_path: str | None = None, ) -> str: st = state_ops.store(config_path) _validate_device(cfg) specs = _input_specs(cfg) - - model: Any - model_artifact: str | None = None - has_explicit_source = bool(from_job or model_s3_uri or onnx_path) - if model_id: - model = model_id - elif prefer_quantized and not has_explicit_source and st.get_last_quantized_model_id(): - model = st.get_last_quantized_model_id() - else: - try: - resolved = resolve_onnx( - cfg=cfg, - output_dir=cfg.aihub.output_dir, - from_job=from_job, - model_s3_uri=model_s3_uri, - onnx_path=onnx_path, - last_training_job=st.get_last_training_job(), - ) - except (FileNotFoundError, ValueError) as e: - CONSOLE.print(f"[red]{e}[/red]") - raise typer.Exit(1) - model = resolved.onnx_path - model_artifact = resolved.model_artifact + try: + source = _resolve_model_source( + cfg, + config_path, + model_id=model_id, + previous_model_id=st.get_last_quantized_model_id(), + from_job=from_job, + model_s3_uri=model_s3_uri, + onnx_path=onnx_path, + ) + except (FileNotFoundError, ValueError) as e: + CONSOLE.print(f"[red]{e}[/red]") + raise typer.Exit(1) try: + hub_model = ( + hub.upload_model(str(source.model), name=cfg.aihub.model_name) + if isinstance(source.model, Path) + else hub.get_model(source.model) + ) result = aihub_jobs.submit_compile_job( - model=model, + model=hub_model, device=cfg.aihub.device, input_specs=specs, target_runtime=cfg.aihub.target_runtime, options=cfg.aihub.compile_options, job_name=_job_name(cfg, "compile"), - model_name=cfg.aihub.model_name if isinstance(model, Path) else None, ) except Exception as e: CONSOLE.print(f"[red]AI Hub compile failed: {e}[/red]") @@ -229,8 +335,8 @@ def _compile_step( "last_compile_job_id": result["job_id"], "last_compiled_model_id": result["model_id"], } - if model_artifact: - updates["last_model_artifact"] = model_artifact + if source.model_artifact: + updates["last_model_artifact"] = source.model_artifact st.update(**updates) CONSOLE.print(f"[green]✓[/green] Compile job: [bold]{result['job_id']}[/bold]") CONSOLE.print(f"[green]✓[/green] Compiled model: [bold]{result['model_id']}[/bold]") @@ -256,8 +362,9 @@ def _validate_step( run = datetime.now().strftime("%Y%m%d-%H%M%S") out_dir = Path(cfg.aihub.output_dir) / run / "validation" try: + hub_model = hub.get_model(resolved_model_id) result = aihub_jobs.submit_inference_job( - resolved_model_id, + hub_model, cfg.aihub.device, inputs, out_dir, @@ -281,8 +388,9 @@ def _profile_step(cfg: Config, config_path: str, model_id: str | None) -> str: _validate_device(cfg) resolved_model_id = _model_id_or_state(config_path, model_id) try: + hub_model = hub.get_model(resolved_model_id) result = aihub_jobs.submit_profile_job( - resolved_model_id, + hub_model, cfg.aihub.device, cfg.aihub.profile_options, job_name=_job_name(cfg, "profile"), @@ -295,9 +403,24 @@ def _profile_step(cfg: Config, config_path: str, model_id: str | None) -> str: return str(result["job_id"]) +@app.command() +def optimize( + from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should optimize"), + model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to optimize"), + onnx_path: str | None = typer.Option( + None, "--onnx-path", help="Local ONNX path or ONNX path inside extracted artifact" + ), + config: str = CONFIG_OPT, +) -> None: + """Optimize an external model into a Workbench-produced ONNX model.""" + cfg = load_cfg(config) + _optimize_step(cfg, config, from_job, model_s3_uri, onnx_path) + + @app.command() def quantize( calibration_path: Path = typer.Argument(..., help="Calibration .npz file or directory of .npy samples"), + model_id: str | None = typer.Option(None, "--model-id", help="AI Hub optimized ONNX model ID"), from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should quantize"), model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to quantize"), onnx_path: str | None = typer.Option( @@ -307,7 +430,15 @@ def quantize( ) -> None: """Quantize an ONNX model to INT8.""" cfg = load_cfg(config) - _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path) + _quantize_step( + cfg, + config, + calibration_path, + model_id=model_id, + from_job=from_job, + model_s3_uri=model_s3_uri, + onnx_path=onnx_path, + ) @app.command() @@ -322,7 +453,14 @@ def compile( ) -> None: """Compile a model for the configured Qualcomm AI Hub target.""" cfg = load_cfg(config) - _compile_step(cfg, config, model_id, from_job, model_s3_uri, onnx_path, prefer_quantized=True) + _compile_step( + cfg, + config, + model_id=model_id, + from_job=from_job, + model_s3_uri=model_s3_uri, + onnx_path=onnx_path, + ) @app.command() @@ -351,7 +489,7 @@ def profile( def upload( calibration_path: Path = typer.Argument(..., help="Calibration .npz file or directory of .npy samples"), input_file: Path = typer.Argument(..., help="Validation .npz or .npy inputs to run on device"), - from_step: UploadStep = typer.Option(UploadStep.quantize, "--from-step", help="Resume from this Workbench step"), + from_step: UploadStep = typer.Option(UploadStep.optimize, "--from-step", help="Resume from this Workbench step"), from_job: str | None = typer.Option(None, "--from-job", help="Training job name whose model artifact should upload"), model_s3_uri: str | None = typer.Option(None, "--model-s3-uri", help="S3 URI of model.tar.gz to upload"), onnx_path: str | None = typer.Option( @@ -360,25 +498,48 @@ def upload( input_name: str | None = typer.Option(None, "--input-name", help="Input name for .npy validation files"), config: str = CONFIG_OPT, ) -> None: - """Run the four Workbench upload steps: quantize, compile, validate, and profile.""" + """Optimize, quantize, optionally compile, validate, and profile a model.""" cfg = load_cfg(config) - steps = [UploadStep.quantize, UploadStep.compile, UploadStep.validate, UploadStep.profile] + steps = [UploadStep.optimize, UploadStep.quantize, UploadStep.compile, UploadStep.validate, UploadStep.profile] selected = steps[steps.index(from_step) :] + optimized_model_id: str | None = None quantized_model_id: str | None = None compiled_model_id: str | None = None + if UploadStep.optimize in selected: + optimized_model_id = _optimize_step(cfg, config, from_job, model_s3_uri, onnx_path) if UploadStep.quantize in selected: - quantized_model_id = _quantize_step(cfg, config, calibration_path, from_job, model_s3_uri, onnx_path) - if UploadStep.compile in selected: - compiled_model_id = _compile_step( + if UploadStep.optimize not in selected: + optimized_model_id = state_ops.store(config).get_last_optimized_model_id() + if not optimized_model_id: + CONSOLE.print( + "[red]No optimized ONNX model found. Resume from --from-step optimize or run " + "'qc-cli ai-hub optimize' first.[/red]" + ) + raise typer.Exit(1) + quantized_model_id = _quantize_step( cfg, config, - model_id=quantized_model_id, - from_job=from_job, - model_s3_uri=model_s3_uri, - onnx_path=onnx_path, - prefer_quantized=True, + calibration_path, + model_id=optimized_model_id, ) + if UploadStep.compile in selected: + if cfg.aihub.target_runtime == "onnx": + compiled_model_id = quantized_model_id or state_ops.store(config).get_last_quantized_model_id() + if not compiled_model_id: + CONSOLE.print( + "[red]No quantized ONNX model found. Resume from --from-step quantize or run " + "'qc-cli ai-hub quantize' first.[/red]" + ) + raise typer.Exit(1) + state_ops.store(config).update(last_compiled_model_id=compiled_model_id) + CONSOLE.print("[green]✓[/green] Target runtime is ONNX; skipping final compile.") + else: + compiled_model_id = _compile_step( + cfg, + config, + model_id=quantized_model_id, + ) if UploadStep.validate in selected: _validate_step(cfg, config, input_file, compiled_model_id, input_name) if UploadStep.profile in selected: diff --git a/src/qualcomm/aihub_jobs.py b/src/qualcomm/aihub_jobs.py index 6afda49..c2c4e9c 100644 --- a/src/qualcomm/aihub_jobs.py +++ b/src/qualcomm/aihub_jobs.py @@ -28,30 +28,19 @@ def _dataset_entries(inputs: dict[str, Any]) -> dict[str, list[Any]]: def submit_compile_job( - model: Any, + model: Model, device: Device, input_specs: dict[str, tuple[tuple[int, ...], str]], target_runtime: str, options: str | None = None, job_name: str | None = None, - model_name: str | None = None, ) -> ModelJobResult: compile_options = f"--target_runtime {target_runtime}" if options: compile_options = f"{compile_options} {options}" - model_arg = model - if isinstance(model, Path): - model_arg = str(model) - elif isinstance(model, str): - candidate = Path(model) - model_arg = model if candidate.exists() or candidate.suffix else hub.get_model(model) - - if model_name and isinstance(model_arg, str) and Path(model_arg).exists(): - model_arg = hub.upload_model(model_arg, name=model_name) - job = hub.submit_compile_job( - model=model_arg, + model=model, device=device, name=job_name, input_specs=input_specs, @@ -64,14 +53,14 @@ def submit_compile_job( def submit_inference_job( - model_id: str, + model: Model, device: Device, inputs: dict[str, Any], output_dir: str | Path, job_name: str | None = None, ) -> InferenceJobResult: job = hub.submit_inference_job( - model=hub.get_model(model_id), + model=model, device=device, inputs=_dataset_entries(inputs), name=job_name, @@ -83,13 +72,13 @@ def submit_inference_job( def submit_profile_job( - model_id: str, + model: Model, device: Device, options: str | None = None, job_name: str | None = None, ) -> ProfileJobResult: job = hub.submit_profile_job( - model=hub.get_model(model_id), + model=model, device=device, name=job_name, options=options or "", @@ -98,17 +87,13 @@ def submit_profile_job( def submit_quantize_job( - model: str | Path, + model: Model, calibration_data: dict[str, Any], options: str | None = None, job_name: str | None = None, - model_name: str | None = None, ) -> ModelJobResult: - model_arg = str(model) - if model_name and Path(model_arg).exists(): - model_arg = hub.upload_model(model_arg, name=model_name) job = hub.submit_quantize_job( - model=model_arg, + model=model, calibration_data=_dataset_entries(calibration_data), weights_dtype=QuantizeDtype.INT8, activations_dtype=QuantizeDtype.INT8, diff --git a/src/state.py b/src/state.py index 3bfcbcb..c273d50 100644 --- a/src/state.py +++ b/src/state.py @@ -37,6 +37,10 @@ class CliStateStore: value = self.get("last_model_artifact") return str(value) if value else None + def get_last_optimized_model_id(self) -> str | None: + value = self.get("last_optimized_model_id") + return str(value) if value else None + def get_last_quantized_model_id(self) -> str | None: value = self.get("last_quantized_model_id") return str(value) if value else None