YOLO Model Training

Frame-by-frame object detection training using Ultralytics YOLO models. Wraps the Ultralytics training API with a fixed interface and CLI entry point, allowing models to be trained from the command line or imported directly into a pipeline.

---

Summary

As with preprocessing, we specify two workflows for training models: local and Sockeye (HPC). Local workflows run linearly, where each training process must be called individually. HPC workflows parallelize training jobs to run each of the eight experimental splits simultaneously.

Property	Detail
Script	training_yolo.py
Framework	Ultralytics YOLO
Supported Models	YOLOv8, YOLOv11 (YOLO26+)
Input	Dataset YAML + pretrained weights / Dataset.tar + pretrained weights
Output	Model weights + training artifacts written to disk
CLI Support	Yes

See Also: 04_train_yolo.sh for the YOLO training Sockeye job submission script and run_training_pipeline.sh for the Sockeye orchestrator script.

---

Inputs

• Local: Reads training data directly from the local dataset.yaml path, loads the pretrained YOLO weights, updates the path references, and runs training.
• Sockeye (HPC): Unpacks a compressed dataset.tar file locally to the compute node to avoid numerous slow file transfers over the HPC network. It updates yaml_path to the current node directory, loads pre-downloaded YOLO weights, and runs training against the freshly unpacked files.

Output

This script writes directly to disk and returns None. All outputs are organized by experimental split and written to project/name/split.

project/
└── name/
    ├── weights/
    │   ├── best.pt       # Best checkpoint evaluated by validation metric
    │   └── last.pt       # Final checkpoint from the last epoch
    ├── results.csv       # Per-epoch training and validation metrics
    └── args.yaml         # Resolved runtime training configuration

---

How It Works

1. Environment Detection: The script detects whether the execution context is local or on Sockeye.

2. On Sockeye, the dataset.tar file is unpacked directly on the local compute node, and the yaml_path variable inside dataset.yaml is updated dynamically to target the node location.

3. On a Local machine, yaml_path inside dataset.yaml is updated to your current project path before training execution.

4. Weight Allocation:

5. On Sockeye, a pretrained YOLO model is loaded from a local, pre-downloaded weights directory (configurable via .env_sockeye and prepared by running bash 01_setup.sh). Ensure the model version downloaded and the variables specified in 04_train_yolo.sh match before running.

6. Locally, weights are fetched directly from Ultralytics using the version and size flags passed via the CLI (e.g., resolving to yolov8n.pt or yolo11n.pt).

7. Training Execution: The model is trained on the dataset described in the generated YAML file.

8. Parameter Forwarding: Core training arguments—such as epochs, batch size, image size, patience, and device constraints—are passed directly to the underlying Ultralytics model.train() execution call.

---

Core Pipeline Functionality

YAML Path Configuration & Dataset Unpacking

Because absolute paths on shared HPC clusters change depending on the assigned compute node, YAML target paths are initially specified as relative links (./dataset). This script resolves and updates paths dynamically at runtime to prevent pathing failures during training.

On Sockeye, datasets are bundled into a single dataset.tar archive to maximize filesystem network efficiency. This script unpacks the dataset onto the local compute node, updates the underlying configuration paths, and validates folder lengths and file counts to ensure complete data integrity before initializing the model.

Image Handling

Training images are dynamically resized to the target img_size resolution. When rect=True (default behavior), the longest side of the frame is restricted to img_size while fully preserving the original aspect ratio. This configuration is highly recommended for high-resolution footage like the 1920×1800 MooVision clips utilized in this pipeline.

Early Stopping

If the chosen validation metrics fail to improve for patience consecutive epochs, training terminates automatically to save compute cycles. You can also pass the time parameter to hard-cap the absolute wall-clock runtime in hours, regardless of how many epochs have finished.

Batch Size

Supports flexible, variable batch configuration. The default batch size on Sockeye is optimized at 64 to maximize throughput; allocating exceptionally large batch values on certain compute nodes may trigger Out Of Memory (OOM) faults.

Epoch Setting

The epochs parameter specifies the maximum number of full training passes over the dataset. While higher epoch caps allow for deep feature optimization, they significantly increase the probability of model overfitting.

Cache Optimization

Allows caching of dataset images and labels directly into memory using --cache. Caching is generally discouraged for this pipeline as the video frame datasets are exceptionally large and risk flooding memory limits. Given the high-speed local disk nodes on Sockeye, caching yields negligible performance benefits.

Overwriting Existing Runs

To prevent accidental data loss, the pipeline prevents overwriting existing directories by default. To replace an existing run folder with a new experiment, pass the --exist_ok flag.

Hardware & GPU Compatibility

The architecture is designed to scale across single or multi-GPU environments using the device and workers parameters. By default, device is set to "cpu" to ensure safe, error-free execution on machines lacking dedicated acceleration.

• To target a specific GPU, pass a single string identifier: --device="0"
• To run parallelized multi-GPU training, pass a collection list: --device=[0,1]
• The workers argument (defaults to 8) regulates the number of parallel CPU data-loading threads used to feed data to the GPU.

Model Selection

Pretrained base weights are dynamically resolved at execution as yolo{version}{size}.pt. For example, setting model=26 (MooVision internally aliases 26 for v11 architectures) and model_size="m" will resolve to yolo11m.pt. Ultralytics automatically handles fetching remote files if matching weights are not discovered locally.

---

Command Line Arguments

For complete information regarding downstream engine configurations, visit the Ultralytics Training Reference Documentation.

Argument	Type	Default	Required	Description
--dataset	str	None	Yes	Path to the dataset configuration target file.
--device	str	"cpu"	No	Hardware acceleration choice. Examples: '0' or '0,1' for specific GPUs, 'cpu', 'cuda', or 'mps'.
--workers	int	8	No	Number of parallel data loader worker processes.
--name	str	"cross-sucking"	No	Name identifier for the specific training run/experiment.
--project	str	"MooVision"	No	Name of the output project directory where results are saved.
--weights_dir	str	""	No	Path to a local folder containing model weights. Leave completely blank to automatically download weights from the internet.
--model	int	26	No	YOLO model version integer (e.g., 8 for YOLOv8, 26 for YOLO11/26+).
--model_size	str	"n"	No	Model scale/variant size configuration options: n (nano), s (small), m (medium), l (large), or x (xlarge).
--epochs	int	50	No	Maximum number of full validation passes over the entire dataframe during training.
--time	float	None	No	Maximum absolute training duration allowed, specified in hours.
--batch	int	32	No	Training batch size per iteration.
--patience	int	20	No	Early stopping threshold. Training stops early if no evaluation improvements are seen after this many consecutive epochs.
--img_size	int	640	No	Target image size (resolution pixel width/height) used for model training.
--not_rect	Flag	True	No	Disables rectangular training. Including this flag switches the underlying configuration (rect) to False, forcing the system to stop maintaining native image aspect ratios.
--do_not_save	Flag	True	No	Disables tracking saves. Including this flag switches the underlying configuration (save) to False, preventing training checkpoints and finalized model weights from saving to disk.
--exist_ok	Flag	False	No	Enables project overwriting. Including this flag sets the parameter to True, allowing the script to safely overwrite existing project directories without throwing errors.
--cache	Flag	False	No	Enables image caching. Including this flag sets the parameter to True, caching raw images in RAM/storage memory for significantly quicker batch step testing.

---

Usage

Command Line (Local) Example

uv run scripts/training/training_yolo.py \
  --dataset="data/training/pipeline_demo/dataset/dataset.yaml" \
  --project="pipeline_demo" \
  --name="demo_01" \
  --device="cpu" \
  --epochs=1 \
  --batch=8

---

Function Reference

Frame-by-frame object detection training using Ultralytics YOLO models.

setup_node_dataset(dataset, base_name='dataset')

Extract and isolate datasets to process-specific local compute node directories.

This function sets up a task-isolated staging workspace inside high-speed node-local storage (e.g., SSD/NVMe scratch spaces mapped via SLURM_TMPDIR) to prevent file collisions when multi-task array jobs run concurrently on the same physical hardware node. It extracts the raw dataset tarball via native system utilities, checks the data structure for YOLO formatting compatibility, and dynamically overwrites the dataset's internal configuration file (dataset.yaml) with absolute pathing strings.

Parameters:

Name	Type	Description	Default
dataset	str	The path string targeting the dataset. On an HPC cluster loop, this maps to the target .tar input file path. On local local dev systems, this maps to a directory relative to ROOT_DIR.	required
base_name	str	An identifier appended to the directory isolation signature naming convention (job_X_task_Y_<base_name>).	"dataset"

Returns:

Name	Type	Description
local_yaml_path	Path	The absolute path pointing to the dynamically updated and isolated configuration dataset.yaml file.
on_cluster	bool	Flags whether execution environments evaluated as an HPC cluster loop or a local fallback run.

Raises:

Type	Description
FileNotFoundError	If the target dataset tar file is missing on a cluster, or if the internal dataset.yaml configuration template cannot be located.
ValueError	If cluster execution conditions are true but SLURM_TMPDIR is missing, or if an image directory component is found completely empty during audit loops.
RuntimeError	If native system tar sub-processes crash or return non-zero exit codes, or if structural validation fails (missing both images and labels trees).

Source code in scripts/training/training_yolo.py

def setup_node_dataset(dataset: str, base_name: str = "dataset") -> Path:
    """
    Extract and isolate datasets to process-specific local compute node directories.

    This function sets up a task-isolated staging workspace inside high-speed
    node-local storage (e.g., SSD/NVMe scratch spaces mapped via `SLURM_TMPDIR`)
    to prevent file collisions when multi-task array jobs run concurrently on the
    same physical hardware node. It extracts the raw dataset tarball via native system
    utilities, checks the data structure for YOLO formatting compatibility, and
    dynamically overwrites the dataset's internal configuration file (`dataset.yaml`)
    with absolute pathing strings.

    Parameters
    ----------
    dataset : str
        The path string targeting the dataset. On an HPC cluster loop, this maps
        to the target `.tar` input file path. On local local dev systems, this maps
        to a directory relative to `ROOT_DIR`.
    base_name : str, default "dataset"
        An identifier appended to the directory isolation signature naming convention
        (`job_X_task_Y_<base_name>`).

    Returns
    -------
    local_yaml_path : pathlib.Path
        The absolute path pointing to the dynamically updated and isolated configuration
        `dataset.yaml` file.
    on_cluster : bool
        Flags whether execution environments evaluated as an HPC cluster loop or a
        local fallback run.

    Raises
    ------
    FileNotFoundError
        If the target dataset tar file is missing on a cluster, or if the internal
        `dataset.yaml` configuration template cannot be located.
    ValueError
        If cluster execution conditions are true but `SLURM_TMPDIR` is missing,
        or if an image directory component is found completely empty during audit loops.
    RuntimeError
        If native system `tar` sub-processes crash or return non-zero exit codes,
        or if structural validation fails (missing both `images` and `labels` trees).
    """
    # GENERATE AN ABSOLUTE ISOLATION SIGNATURE PER ARRAY TASK
    # This ensures that even if Task 1 and Task 2 share the same physical server node,
    # they write to completely separated folders on the NVMe storage drive.
    slurm_job_id = os.environ.get("SLURM_JOB_ID", "local_dev")
    slurm_task_id = os.environ.get("SLURM_ARRAY_TASK_ID", "0")
    tmp_dir_env = os.environ.get("SLURM_TMPDIR")

    on_cluster = "PBS_JOBID" in os.environ or "SLURM_JOB_ID" in os.environ

    if on_cluster:

        print("Working on Cluster, searching for tar path...")
        tar_path = Path(dataset)
        if not tar_path.exists():
            raise FileNotFoundError(f"Dataset archive missing at: {tar_path}")

        if tmp_dir_env:
            # Combined both Job ID and Task ID to guarantee absolute isolation
            # across different jobs and task arrays sharing the same node.
            # e.g., /localscratch/11740177/job_11740177_task_1_dataset
            isolated_node_dir = (
                Path(tmp_dir_env)
                / f"job_{slurm_job_id}_task_{slurm_task_id}_{base_name}"
            )
        else:
            raise ValueError("Could not find SLURM_TMPDIR.")

        # ISOLATED NATIVE C TAR EXTRACTION
        # If this specific task has already successfully extracted its partition, skip to avoid overhead
        if not (isolated_node_dir.exists() and any(isolated_node_dir.iterdir())):
            isolated_node_dir.mkdir(parents=True, exist_ok=True)
            print(f"[INFO] Task {slurm_task_id} Staging Area: {isolated_node_dir}")
            print(f"[INFO] Extracting archive via native system tar utility...")

            try:
                # --strip-components=1 strips the top-level folder 'dataset/' from the tarball
                # and unrolls everything straight into our task-isolated directory
                subprocess.run(
                    [
                        "tar",
                        "-xf",
                        str(tar_path),
                        "--strip-components=1",
                        "-C",
                        str(isolated_node_dir),
                    ],
                    check=True,
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.PIPE,
                )
            except subprocess.CalledProcessError as e:
                if isolated_node_dir.exists():
                    shutil.rmtree(isolated_node_dir)
                error_msg = (
                    e.stderr.decode().strip() if e.stderr else "I/O System Error"
                )
                raise RuntimeError(
                    f"[ERROR] Native tar extraction failed for task {slurm_task_id}: {error_msg}"
                )

        # POST-EXTRACTION VALIDATION CHECKS
        print(
            f"[INFO] Task {slurm_task_id}: Commencing dataset integrity validation..."
        )

        # Dynamically check for standard YOLO data splits: images and labels
        valid_folders = [
            d.name
            for d in isolated_node_dir.iterdir()
            if d.is_dir()
            and not d.name.startswith(".")
            and d.name in ["images", "labels"]
        ]

        # Should return "images" and "labels", else clean for next run.
        if not valid_folders:
            if isolated_node_dir.exists():
                shutil.rmtree(isolated_node_dir)
            raise RuntimeError(
                f"[ERROR] Validation Failed: No valid YOLO folders ('images'/'labels') found inside "
                f"{isolated_node_dir}. Verify if '--strip-components=1' fits your archive structure."
            )

        # Path(tmp_dir_env) / f"job_{slurm_job_id}_task_{slurm_task_id}_{base_name}" / "images" or "labels"
        img_root_dir = isolated_node_dir / valid_folders[0]
        lbl_root_dir = isolated_node_dir / valid_folders[1]

        # Should return "train" and "val"
        valid_splits = [
            s.name
            for s in img_root_dir.iterdir()
            if s.is_dir() and not s.name.startswith(".") and s.name in ["train", "val"]
        ]
        # Step B: Traverse each found data split to audit structural files
        for split in valid_splits:

            # ... / images / train or val
            img_dir = img_root_dir / split
            lbl_dir = lbl_root_dir / split

            if not img_dir.exists() or not lbl_dir.exists():
                if isolated_node_dir.exists():
                    shutil.rmtree(isolated_node_dir)
                raise FileNotFoundError(
                    f"[ERROR] Structural Mismatch in split [{split.upper()}]. "
                    f"Expected both branches to exist:\n  - {img_dir}\n  - {lbl_dir}"
                )

            # Map unique stems (filenames without extensions) while dropping hidden OS assets
            image_stems = {
                f.stem
                for f in img_dir.iterdir()
                if f.is_file() and not f.name.startswith(".")
            }
            label_stems = {
                f.stem
                for f in lbl_dir.iterdir()
                if f.is_file() and not f.name.startswith(".")
            }

            num_images = len(image_stems)
            num_labels = len(label_stems)

            # Step C: Catch empty directory states (indicates quota block or corrupted extraction)
            if num_images == 0:
                if isolated_node_dir.exists():
                    shutil.rmtree(isolated_node_dir)
                raise ValueError(
                    f"[ERROR] Integrity Failure: Image directory for split [{split.upper()}] is completely empty."
                )

            # Step D: Alert or flag on image vs bounding-box count variance
            if num_images != num_labels:
                print(
                    f"[WARNING] File count discrepancy in split [{split.upper()}]: "
                    f"Found {num_images} images but {num_labels} labels."
                )

                # Check for images missing corresponding annotation labels
                orphaned_images = image_stems - label_stems
                if orphaned_images:
                    print(
                        f"First 5 images missing a matching .txt label file: {list(orphaned_images)[:5]}"
                    )

            print(
                f"[SUCCESS] Split [{split.upper()}] checked: {num_images} images and {num_labels} labels verified."
            )

        print(
            f"[SUCCESS] Dataset structure validation completely passed for Task {slurm_task_id}."
        )
        # ────────────────────────────────────────────────────────────────
    else:
        # Laptop fallback strategy
        data_path = Path(ROOT_DIR) / dataset
        isolated_node_dir = data_path.parent

    # CONFIGURE INDEPENDENT PATHS IN THE LOCAL YAML FILE
    # The config file sits inside each task's isolated local folder footprint,
    # meaning tasks will never read or overwrite each other's paths.
    local_yaml_path = isolated_node_dir / "dataset.yaml"
    if not local_yaml_path.exists():
        raise FileNotFoundError(
            f"Missing dataset.yaml configuration file inside task directory: {isolated_node_dir}"
        )

    # Read the isolated config file
    with open(local_yaml_path, "r") as f:
        config_data = yaml.safe_load(f)

    # Re-route the 'path' key from './dataset' to the absolute path of this task's folder
    config_data["path"] = str(isolated_node_dir.resolve())

    # Save the modifications back to the node-local storage drive
    with open(local_yaml_path, "w") as f:
        yaml.dump(config_data, f, default_flow_style=False, sort_keys=False)

    if on_cluster:
        print(
            f"[SUCCESS] Task {slurm_task_id} absolute dataset path locked to: {config_data['path']}"
        )
    else:
        print(f"[SUCCESS] Absolute dataset path locked to: {config_data['path']}")
    return local_yaml_path, on_cluster

train_yolo_model(on_cluster, yaml_path, name, project, weights_dir, model, model_size, device, workers, exist_ok=False, epochs=100, time=None, patience=50, batch=16, img_size=640, save=True, rect=True, cache=False, **kwargs)

Train a YOLO model.

Takes in training set from preprocessing.py and trains a YOLOv8 model of the specified size.

Parameters are taken from ultralytics docs. For more information and arguments see: https://docs.ultralytics.com/modes/train.

Parameters:

Name	Type	Description	Default
yaml_path	str	Path to data.yaml configuration file. File contains dataset paths, and classes (e.g. cross-sucking)	required
name	str	Experiment name. e.g. cross-sucking-yolo-v26-m	required
project	Path | str	Output directory. This is where output of model training is stored, including model weights and predictions.	required
weights_dir	str	Path to local pretrained YOLO weights (raw weights)	required
model	int	YOLO model version. i.e. 26 for v26, 8 for v8.	required
model_size	str	Model size (n=nano, s=small, m=medium, l=large, x=xlarge). A larger model may take longer to train or use, but it may also provide increased performance.	required
device	str	device to train on (0 for GPU, 'cuda' for CUDA, 'cpu' for CPU, 'mps' for GPU on Mac).	required
workers	int	Number of parallel threads to run on sockeye.	required
exist_ok	bool	If True, overwrite existing project/name directory.	False
epochs	int	Number of training epochs. Each epoch represents one full pass of the dataest. A larger epoch value will increase training duration, but may also increase accuracy.	100
time	float	Maximum training time in hours. This will override the number of epochs if reached first. Note, this may affect model performance and training times.	None
patience	int	Number of epochs to wait until stopping the training due no improvements in validation metrics.	50
batch	int or float	Batch size. See ultrlytics docs for more info: https://www.ultralytics.com/glossary/batch-size	16
img_size	int	Image size for training. Resized to squares of img_size x img_size if rect = False, else keeps aspect ratio. This may affect model accuracy and computational complexity.	640
save	bool	Enables saving of training checkpoints and final model weights. Set to True by default to save results of model training.	True
rect	bool	If True, keeps image ratio when rescaling images. The longest side is equal to img_size. Can improve efficiency and speed but may affect accuracy.	True
**kwargs	named arguemnts	See ultralytics docs for additional arguments.	{}

Returns:

Type	Description
None	This function reads to disk and does not return anything.

Examples:

.. code-block:: python

# After running scripts read_all_clips_index.py, and splitting.py run...
yaml_path = "data/processed/pipeline_testing/yolo_format"  #/path/to/yaml/file
train_yolo_model(
    yaml_path=yaml_path,
    model_size="n",
    epochs=3,
    imgsz=640,
    batch=8,
    rect=True,
    device="mps",
    save=True,
    patience=50,
    name="cross_sucking",
    project="MooVision",
)

Source code in scripts/training/training_yolo.py

def train_yolo_model(
    on_cluster: bool,
    yaml_path: str,
    name: str,
    project: Path | str,
    weights_dir: str,
    model: int,
    model_size: int,
    device: str,
    workers: int,
    exist_ok: bool = False,
    epochs: int = 100,
    time: float = None,
    patience: int = 50,
    batch: int | float = 16,
    img_size: int = 640,
    save: bool = True,
    rect: bool = True,
    cache=False,
    **kwargs,
):
    """
    Train a YOLO model.

    Takes in training set from preprocessing.py and trains a YOLOv8 model of
    the specified size.

    Parameters are taken from ultralytics docs. For more information and
    arguments see:
    https://docs.ultralytics.com/modes/train.


    Parameters
    ----------
    yaml_path : str
        Path to data.yaml configuration file. File contains dataset paths,
        and classes (e.g. cross-sucking)
    name : str
        Experiment name. e.g. cross-sucking-yolo-v26-m
    project : Path | str
        Output directory. This is where output of model training is stored, including model
        weights and predictions.
    weights_dir : str
        Path to local pretrained YOLO weights (raw weights)
    model : int
        YOLO model version. i.e. 26 for v26, 8 for v8.
    model_size : str
        Model size (n=nano, s=small, m=medium, l=large, x=xlarge). A larger
        model may take longer to train or use, but it may also provide
        increased performance.
    device : str
        device to train on (0 for GPU, 'cuda' for CUDA, 'cpu' for CPU, 'mps' for GPU on Mac).
    workers : int
        Number of parallel threads to run on sockeye.
    exist_ok : bool
        If True, overwrite existing project/name directory.
    epochs : int
        Number of training epochs. Each epoch represents one full pass of the
        dataest. A larger epoch value will increase training duration, but may
        also increase accuracy.
    time : float
        Maximum training time in hours. This will override the number of epochs
        if reached first. Note, this may affect model performance and training
        times.
    patience : int
        Number of epochs to wait until stopping the training due no
        improvements in validation metrics.
    batch : int or float
        Batch size. See ultrlytics docs for more info:
        https://www.ultralytics.com/glossary/batch-size
    img_size : int
        Image size for training. Resized to squares of img_size x img_size if
        rect = False, else keeps aspect ratio. This may affect model accuracy
        and computational complexity.
    save : bool
        Enables saving of training checkpoints and final model weights. Set to True by default
        to save results of model training.
    rect : bool
        If True, keeps image ratio when rescaling images. The longest side is
        equal to img_size. Can improve efficiency and speed but may affect
        accuracy.
    **kwargs : named arguemnts
        See ultralytics docs for additional arguments.


    Returns
    -------
    None
        This function reads to disk and does not return anything.


    Examples
    --------
    .. code-block:: python

        # After running scripts read_all_clips_index.py, and splitting.py run...
        yaml_path = "data/processed/pipeline_testing/yolo_format"  #/path/to/yaml/file
        train_yolo_model(
            yaml_path=yaml_path,
            model_size="n",
            epochs=3,
            imgsz=640,
            batch=8,
            rect=True,
            device="mps",
            save=True,
            patience=50,
            name="cross_sucking",
            project="MooVision",
        )
    """
    # Loads pretrained model
    if weights_dir:
        print("Loading model from weights...")
        if model == 26:
            print("Option 1: Loading YOLO26 model")
            final_model_target = Path(weights_dir) / f"yolo{model}{model_size}.pt"
            if not final_model_target.exists():
                raise FileNotFoundError(
                    f"Could not find {final_model_target}. Set model to yolo{model}{model_size}.pt in .env_sockeye and run ./01_setup.sh on the login node."
                )
        else:
            final_model_target = Path(weights_dir) / f"yolov{model}{model_size}.pt"
            if not final_model_target.exists():
                raise FileNotFoundError(
                    f"Could not find {final_model_target}. Set model to yolov{model}{model_size}.pt in .env_sockeye and run ./01_setup.sh on the login node."
                )
    else:
        print("loading model from internet...")
        if model == 26:
            final_model_target = f"yolo{model}{model_size}.pt"

        else:
            final_model_target = f"yolov{model}{model_size}.pt"

    model = YOLO(final_model_target)
    print(f"Model loaded from: {final_model_target}")

    if on_cluster:
        # Path to save in root dir defined in sockeye scripts.
        model.train(
            data=yaml_path,
            name=name,
            project=project,
            device=device,
            exist_ok=exist_ok,
            epochs=epochs,
            time=time,
            patience=patience,
            batch=batch,
            imgsz=img_size,
            save=save,
            rect=rect,
            workers=workers,
            cache=cache,
            **kwargs,
        )
    else:
        # Save project to root dir
        model.train(
            data=yaml_path,
            name=name,
            project=Path(ROOT_DIR) / "data" / "yolo_training_runs" / project,
            device=device,
            exist_ok=exist_ok,
            epochs=epochs,
            time=time,
            patience=patience,
            batch=batch,
            imgsz=img_size,
            save=save,
            rect=rect,
            workers=workers,
            cache=cache,
            **kwargs,
        )

options: show_source: false show_root_heading: true

---

Work In Progress (WIP)

1. Flexible Parameter Forwarding: Integrate kwargs capability into parsed argument behaviors. This will unlock direct passthrough optimization, allowing users to leverage any latent arguments supported natively within the Ultralytics API straight from this interface wrapper.

Target Interface Concept

Pass additional Ultralytics arguments via --kwargs (feature implementation pending):

python train_yolo_model.py --yaml_path dataset.yaml --kwargs lr0=0.01 cos_lr=True