131 lines
4.5 KiB
Python
Executable file
131 lines
4.5 KiB
Python
Executable file
"""Run one small CPU training epoch for every task supported by the WebUI."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
import traceback
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from yolo_webui import TrainingConfig, TrainingRunner
|
|
from yolo_webui.config import AugmentationConfig, MlflowConfig
|
|
|
|
|
|
TASKS = {
|
|
"detect": ("datasets/yolo26_smoke/detect/detect.yaml", "yolo26n.pt"),
|
|
"segment": ("datasets/yolo26_smoke/segment/segment.yaml", "yolo26n-seg.pt"),
|
|
"classify": ("datasets/yolo26_smoke/classify", "yolo26n-cls.pt"),
|
|
"pose": ("datasets/yolo26_smoke/pose/pose.yaml", "yolo26n-pose.pt"),
|
|
"obb": ("datasets/yolo26_smoke/obb/obb.yaml", "yolo26n-obb.pt"),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--mlflow",
|
|
action="store_true",
|
|
help="Enable MLflow logging for every smoke-training run.",
|
|
)
|
|
parser.add_argument(
|
|
"--tracking-uri",
|
|
default="sqlite:///runs/yolo26_mlflow_smoke/mlflow.db",
|
|
)
|
|
parser.add_argument("--experiment", default="yolo26-mlflow-smoke")
|
|
parser.add_argument("--project", type=Path)
|
|
args = parser.parse_args()
|
|
|
|
results: dict[str, dict[str, object]] = {}
|
|
default_project = "runs/yolo26_mlflow_smoke/train" if args.mlflow else "runs/yolo26_smoke"
|
|
project_dir = (args.project or Path(default_project)).resolve()
|
|
smoke_batch_id = uuid.uuid4().hex if args.mlflow else ""
|
|
if args.mlflow:
|
|
import mlflow
|
|
|
|
project_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
if args.tracking_uri.startswith("sqlite:///"):
|
|
tracking_db = Path(args.tracking_uri.removeprefix("sqlite:///"))
|
|
tracking_db.parent.mkdir(parents=True, exist_ok=True)
|
|
os.environ["YOLO_WEBUI_MLFLOW_RUN_GROUP"] = smoke_batch_id
|
|
mlflow.set_tracking_uri(args.tracking_uri)
|
|
mlflow.set_experiment(args.experiment)
|
|
with mlflow.start_run(run_name=f"smoke-batch-{smoke_batch_id}") as anchor:
|
|
mlflow.set_tags(
|
|
{
|
|
"smoke.anchor": "true",
|
|
"yolo.run_group": smoke_batch_id,
|
|
}
|
|
)
|
|
print(f"MLflow smoke batch: {smoke_batch_id}", flush=True)
|
|
|
|
for task, (dataset, model) in TASKS.items():
|
|
print(f"\n=== {task}: {model} ===", flush=True)
|
|
config = TrainingConfig(
|
|
dataset=dataset,
|
|
model=model,
|
|
task=task, # type: ignore[arg-type]
|
|
epochs=1,
|
|
image_size=64,
|
|
batch_size=2,
|
|
device="cpu",
|
|
workers=0,
|
|
patience=0,
|
|
project=str(project_dir),
|
|
run_name=task,
|
|
augmentation=AugmentationConfig(enabled=False),
|
|
mlflow=MlflowConfig(
|
|
enabled=args.mlflow,
|
|
tracking_uri=args.tracking_uri,
|
|
experiment_name=args.experiment,
|
|
run_name=f"{task}-smoke" if args.mlflow else "",
|
|
),
|
|
)
|
|
started = time.monotonic()
|
|
try:
|
|
output = TrainingRunner().train(
|
|
config,
|
|
lambda event: print(
|
|
f"[{event.kind}] {event.message}",
|
|
flush=True,
|
|
),
|
|
)
|
|
results[task] = {
|
|
"status": "succeeded",
|
|
"seconds": round(time.monotonic() - started, 2),
|
|
"output": str(output) if output else None,
|
|
}
|
|
except Exception as exc:
|
|
traceback.print_exc()
|
|
results[task] = {
|
|
"status": "failed",
|
|
"seconds": round(time.monotonic() - started, 2),
|
|
"error": f"{type(exc).__name__}: {exc}",
|
|
}
|
|
|
|
summary_path = (
|
|
project_dir.parent / "smoke_summary.json"
|
|
if args.mlflow
|
|
else project_dir / "smoke_summary.json"
|
|
)
|
|
summary_path.parent.mkdir(parents=True, exist_ok=True)
|
|
summary = {
|
|
"smoke_batch_id": smoke_batch_id or None,
|
|
"tracking_uri": args.tracking_uri if args.mlflow else None,
|
|
"experiment": args.experiment if args.mlflow else None,
|
|
"tasks": results,
|
|
}
|
|
summary_path.write_text(
|
|
json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"\nSummary: {summary_path.resolve()}")
|
|
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
|
if any(result["status"] != "succeeded" for result in results.values()):
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|