149 lines
5.1 KiB
Python
149 lines
5.1 KiB
Python
#!/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()
|