initial version to train yolo model
This commit is contained in:
172
examples/meter-detection/README.md
Normal file
172
examples/meter-detection/README.md
Normal file
@@ -0,0 +1,172 @@
|
||||
# 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
|
||||
|
||||
- AWS credentials configured for the profile in `config.yaml`
|
||||
- Infrastructure already deployed with `uv run qc-cli infra setup`
|
||||
- A Roboflow API key exported as `ROBOFLOW_API_KEY`
|
||||
- `curl` and `unzip` available locally
|
||||
|
||||
Install or sync the project dependencies:
|
||||
|
||||
```bash
|
||||
uv sync
|
||||
```
|
||||
|
||||
Set the Roboflow API key for the current shell:
|
||||
|
||||
```bash
|
||||
export ROBOFLOW_API_KEY=your-roboflow-api-key
|
||||
```
|
||||
|
||||
## 1. Download The Dataset
|
||||
|
||||
Download version 1 of the dataset in YOLO format. The script uses the Roboflow REST API directly and does not require Python:
|
||||
|
||||
```bash
|
||||
bash examples/meter-detection/download_dataset.sh
|
||||
```
|
||||
|
||||
Confirm the extracted 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
|
||||
```
|
||||
|
||||
The expected layout is similar to:
|
||||
|
||||
```text
|
||||
examples/meter-detection/data/electric-meter-detection/
|
||||
data.yaml
|
||||
train/
|
||||
valid/
|
||||
test/
|
||||
```
|
||||
|
||||
The `test/` split may be absent depending on the exported dataset version.
|
||||
|
||||
## 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
|
||||
uv run qc-cli infra status --config config.yaml
|
||||
```
|
||||
|
||||
## 4. Upload The Dataset
|
||||
|
||||
Upload the downloaded Roboflow dataset to the `s3.data_prefix` configured in `config.yaml`:
|
||||
|
||||
```bash
|
||||
uv run qc-cli upload examples/meter-detection/data/electric-meter-detection --config config.yaml
|
||||
```
|
||||
|
||||
Directory uploads preserve paths relative to the uploaded directory, so SageMaker receives the dataset root with `data.yaml` plus the split directories.
|
||||
|
||||
## 5. Start Training
|
||||
|
||||
Submit the SageMaker training job:
|
||||
|
||||
```bash
|
||||
uv run qc-cli train start --config config.yaml
|
||||
```
|
||||
|
||||
The command prints the submitted SageMaker job name. Check progress with:
|
||||
|
||||
```bash
|
||||
uv run qc-cli train status --config config.yaml
|
||||
```
|
||||
|
||||
Or pass the job name explicitly:
|
||||
|
||||
```bash
|
||||
uv run qc-cli train status qc-cli-YYYYMMDD-HHMMSS --config config.yaml
|
||||
```
|
||||
|
||||
## 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`.
|
||||
|
||||
## 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 `SM_CHANNEL_TRAIN`. |
|
||||
|
||||
Do not set `train-dir` or `model-dir` in normal SageMaker runs. SageMaker sets those automatically through `SM_CHANNEL_TRAIN` and `SM_MODEL_DIR`.
|
||||
57
examples/meter-detection/download_dataset.sh
Executable file
57
examples/meter-detection/download_dataset.sh
Executable file
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
WORKSPACE="kemals-workspace-kbc8l"
|
||||
PROJECT="electric-meter-detection-o4tfi"
|
||||
VERSION="1"
|
||||
FORMAT="yolov8"
|
||||
DATASET_DIR="examples/meter-detection/data/electric-meter-detection"
|
||||
|
||||
if [[ -z "${ROBOFLOW_API_KEY:-}" ]]; then
|
||||
echo "ROBOFLOW_API_KEY is required." >&2
|
||||
echo "Run: export ROBOFLOW_API_KEY=your-roboflow-api-key" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! command -v curl >/dev/null 2>&1; then
|
||||
echo "curl is required." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! command -v unzip >/dev/null 2>&1; then
|
||||
echo "unzip is required." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf "${TMP_DIR}"' EXIT
|
||||
|
||||
API_URL="https://api.roboflow.com/${WORKSPACE}/${PROJECT}/${VERSION}/${FORMAT}?api_key=${ROBOFLOW_API_KEY}"
|
||||
RESPONSE_FILE="${TMP_DIR}/roboflow-export.json"
|
||||
ZIP_FILE="${TMP_DIR}/dataset.zip"
|
||||
|
||||
echo "Requesting Roboflow export link..."
|
||||
curl -fsSL "${API_URL}" -o "${RESPONSE_FILE}"
|
||||
|
||||
DOWNLOAD_URL="$(
|
||||
sed -n 's/.*"link"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' "${RESPONSE_FILE}" \
|
||||
| head -n 1 \
|
||||
| sed 's#\\/#/#g; s#\\u0026#\&#g'
|
||||
)"
|
||||
|
||||
if [[ -z "${DOWNLOAD_URL}" ]]; then
|
||||
echo "Could not find export.link in Roboflow response." >&2
|
||||
echo "Response:" >&2
|
||||
cat "${RESPONSE_FILE}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
mkdir -p "${DATASET_DIR}"
|
||||
|
||||
echo "Downloading dataset ZIP..."
|
||||
curl -fL "${DOWNLOAD_URL}" -o "${ZIP_FILE}"
|
||||
|
||||
echo "Extracting dataset..."
|
||||
unzip -q -o "${ZIP_FILE}" -d "${DATASET_DIR}"
|
||||
|
||||
echo "Downloaded dataset to ${DATASET_DIR}"
|
||||
3
examples/meter-detection/source/requirements.txt
Normal file
3
examples/meter-detection/source/requirements.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
ultralytics>=8.3.0
|
||||
pyyaml>=6.0.3
|
||||
onnx>=1.16.0
|
||||
148
examples/meter-detection/source/train.py
Normal file
148
examples/meter-detection/source/train.py
Normal file
@@ -0,0 +1,148 @@
|
||||
#!/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 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("--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 find_data_yaml(train_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(train_dir.rglob("data.yaml"))
|
||||
if not matches:
|
||||
raise FileNotFoundError(f"Could not find data.yaml under {train_dir}")
|
||||
if len(matches) > 1:
|
||||
print(f"Found multiple data.yaml files; using {matches[0]}")
|
||||
return matches[0]
|
||||
|
||||
|
||||
def _split_exists(dataset_root: Path, value: Any) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
split_path = Path(str(value))
|
||||
if split_path.is_absolute():
|
||||
return split_path.exists()
|
||||
return (dataset_root / split_path).exists()
|
||||
|
||||
|
||||
def prepare_data_yaml(data_yaml: Path) -> Path:
|
||||
"""Write a SageMaker-local data file with absolute dataset paths."""
|
||||
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)
|
||||
|
||||
for split_name in ("train", "val", "valid", "test"):
|
||||
split_value = normalized.get(split_name)
|
||||
if split_value is None:
|
||||
continue
|
||||
split_path = Path(str(split_value))
|
||||
if split_path.is_absolute():
|
||||
normalized[split_name] = str(split_path)
|
||||
else:
|
||||
normalized[split_name] = str((dataset_root / split_path).resolve())
|
||||
|
||||
if "val" not in normalized and "valid" in normalized:
|
||||
normalized["val"] = normalized["valid"]
|
||||
|
||||
if not _split_exists(dataset_root, normalized.get("train")):
|
||||
raise FileNotFoundError(f"Could not resolve train split from {data_yaml}")
|
||||
if not _split_exists(dataset_root, normalized.get("val")):
|
||||
raise FileNotFoundError(f"Could not resolve validation split from {data_yaml}")
|
||||
|
||||
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()
|
||||
train_dir = Path(args.train_dir)
|
||||
model_dir = Path(args.model_dir)
|
||||
model_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
data_yaml = prepare_data_yaml(find_data_yaml(train_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))
|
||||
copy_if_exists(onnx_path, model_dir / "model.onnx")
|
||||
|
||||
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(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()
|
||||
Reference in New Issue
Block a user