update

examples/meter-detection/source/train.py

@@ -24,37 +24,29 @@ def parse_args() -> argparse.Namespace:
     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("--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(train_dir: Path, explicit_path: str | None) -> Path:
+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(train_dir.rglob("data.yaml"))
+    matches = sorted(dataset_dir.rglob("data.yaml"))
     if not matches:
-        raise FileNotFoundError(f"Could not find data.yaml under {train_dir}")
+        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 _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."""
+    """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):
@@ -62,24 +54,8 @@ def prepare_data_yaml(data_yaml: Path) -> Path:
 
     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}")
+        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")
@@ -95,11 +71,11 @@ def copy_if_exists(source: Path, destination: Path) -> None:
 
 def main() -> None:
     args = parse_args()
-    train_dir = Path(args.train_dir)
+    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(train_dir, args.data_yaml))
+    data_yaml = prepare_data_yaml(find_data_yaml(dataset_dir, args.data_yaml))
     model = YOLO(args.model)
 
     train_kwargs: dict[str, Any] = {