YOLO Training Pipeline

YOLO Training Pipeline

This guide describes how to prepare a dataset and train a custom YOLO model using annotations created with Label Studio.

Prerequisites

Before starting, ensure that:

• Label Studio has been installed and configured
• A dataset has been annotated and exported
• Python 3.10 or later is available
• Ultralytics YOLO is installed

Extract the Dataset

After annotating the documents in Label Studio (see Label Studio Guide), extract the downloaded archive to a local directory.

The extracted archive contains:

• labels/ (YOLO annotation files)
• classes.txt (class definitions)
• notes.json (Label Studio metadata)
• images/ (may be empty depending on export settings)

YOLO archive content

Prepare the folder to be correctly used for the training

Dataset Structure

The training pipeline expects the following directory structure:

dataset/
├── images/
│   ├── train/
│   └── val/
├── labels/
│   ├── train/
│   └── val/
├── classes.txt
└── notes.json

Where:

• images/train contains the training images
• images/val contains the validation images
• labels/train contains the training annotations
• labels/val contains the validation annotations
• data.yaml contains the dataset configuration

IMPORTANT

Move the .txt file from labels directory into labels/train directory.

And create the other required folders.

IMPORTANT

The conversion process generates only the annotation files and does not move the images.

Copy the corresponding images into the images/train directory.

Do not worry about populating the images/val or labels/val directories.

Those will be automatically populated moving the 30% of files from images/train and labels/train via a script

Verify that every annotation file has a matching image. Each image must have a corresponding annotation file with the same filename.

Example:

images/train/invoice001.jpg
labels/train/invoice001.txt

IMPORTANT

Rename classes.txt to synset.txt, as the training pipeline expects class definitions to be provided through a file named synset.txt.

In addition to that specific structure for the dataset/ directory, other important files are required for the training (data.yaml, train.py).

dataset/
├── images/
│   ├── train/
│   └── val/
├── labels/
│   ├── train/
│   └── val/
├── notes.json
├── data.yaml
├── synset.txt
└── train.py

Create the data.yaml File

YOLO uses a configuration file named data.yaml to locate the dataset and identify the available classes.

The file specifies:

• Training image directory
• Validation image directory
• Number of classes
• Class names

Example:

path: dataset
train: images/train
val: images/val

nc: 4

names:
  0: invoice_number
  1: date
  2: seller_name
  3: total

The class definitions must match those used during annotation.

Create the Python Training Script

The training script shoul include the Ultralytics YOLO framework.

A pretrained YOLO model such as:

• yolo11n.pt
• yolo11s.pt
• yolo11m.pt
• yolo11l.pt

is required as the starting point for training.

The model does not need to be downloaded manually. Ultralytics automatically downloads it when the training script is executed.

A minimal training script can be as simple as:

from ultralytics import YOLO

model = YOLO("yolo11n.pt")

model.train(
    data="data.yaml",
    epochs=100
)

A more specific training script can look like this:

from ultralytics import YOLO

def main():
    model = YOLO("yolo11m.pt")

    model.train(
        data="data.yaml",
        epochs=150,
        batch=15,
        workers=12,
        lr0=0.002,
        momentum=0.9,
        weight_decay=0.0005,
        warmup_epochs=3,
        warmup_momentum=0.8,
        warmup_bias_lr=0.1,
        optimizer="AdamW",
        patience=30,
        imgsz=846,
        name="target",
        device=0,
        amp=True,
        save=True,
        save_period=10,
        cache="disk"
    )

if __name__ == "__main__":
    main()

The parameters most commonly customized are:

• epochs - number of training epochs
• batch - number of images processed per batch
• imgsz - input image size
• device - CPU or GPU device used for training
• project and name - output directory configuration

for a comprehensive list of parameters: https://docs.ultralytics.com/modes/train?utm_source=chatgpt.com#musgd-optimizer

Start the Training

During training, Ultralytics YOLO:

• Loads the dataset
• Trains the model
• Evaluates the model on the validation set
• Saves the best-performing weights

To launch the training and automate the distribution of the images and labels across the train/ and val/ directories,

You can use a train.sh script, or you can manually move the desired amount of images and labels in the val/ directories and launch the train.py script

To work, this train.sh file, should be included as well in the training directory. In addition to that, given that the script reference a convert-pt-to-onnx.py script to perform the conversion.

this script should be added as well. An example of convert-pt-to-onnx.py can be found here: YOLO to ONNX Conversion.

dataset/
├── images/
│   ├── train/
│   └── val/
├── labels/
│   ├── train/
│   └── val/
├── notes.json
├── data.yaml
├── synset.txt
├── train.py
└── train.sh

This is a train.sh file you can use:

#!/bin/bash

IMAGES_TRAINING_DIR=dataset/images/train
IMAGES_VALIDATION_DIR=dataset/images/val
LABELS_TRAINING_DIR=dataset/labels/train
LABELS_VALIDATION_DIR=dataset/labels/val
VALIDATION_RATE=0.0

TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
LOG=logs/training-$TIMESTAMP.log

# Remove .npy files perhaps left by last elaboration
rm -rf dataset/images/train/*.npy
rm -rf dataset/images/val/*.npy



echo "$TIMESTAMP: Looking for unlabeled images in: $IMAGES_TRAINING_DIR" >> $LOG
for file in $IMAGES_TRAINING_DIR/*.*; do
  BASE_FILENAME="$(basename "$file")"
  BASE_FILENAME="${BASE_FILENAME%.*}"        # strip the extension

  if [ ! -f "$LABELS_TRAINING_DIR/$BASE_FILENAME.txt" ]; then
    TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
    echo "$TIMESTAMP: Image $file non annotated, delete it" >> $LOG
    rm -rf $file
  fi
done

TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
echo "$TIMESTAMP: Looking for unlabeled images in: $IMAGES_VALIDATION_DIR" >> $LOG
for file in $IMAGES_VALIDATION_DIR/*.*; do
  BASE_FILENAME="$(basename "$file")"
  BASE_FILENAME="${BASE_FILENAME%.*}"        # strip the extension

  if [ ! -f "$LABELS_VALIDATION_DIR/$BASE_FILENAME.txt" ]; then
    TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
    echo "$TIMESTAMP: Image $file non annotated, delete it" >> $LOG
    rm -rf $file
  fi
done

echo "$TIMESTAMP: Take $VALIDATION_RATE percent of training images and move them to validation" >> $LOG
total=$(ls -1 $IMAGES_TRAINING_DIR | wc -l)
take=$(echo "$total * $VALIDATION_RATE" | bc | cut -d. -f1)
for file in $(ls -1 $IMAGES_TRAINING_DIR | shuf | head -n "$take"); do
  BASE_FILENAME="$(basename "$file")"
  BASE_FILENAME="${BASE_FILENAME%.*}"        # strip the extension

  TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
  echo "$TIMESTAMP: Relocating file: $file" >> $LOG
  mv $IMAGES_TRAINING_DIR/$file  $IMAGES_VALIDATION_DIR
  mv $LABELS_TRAINING_DIR/$BASE_FILENAME.txt $LABELS_VALIDATION_DIR
done

TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
echo "$TIMESTAMP: Launching the training" >> $LOG
python train.py >> $LOG  2>&1
TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
echo "$TIMESTAMP: Training completed" >> $LOG


TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
echo "$TIMESTAMP: Launching the conversion" >> $LOG
python convert_pt_to_onnx.py >> $LOG 2>&1
TIMESTAMP=$(date +"%Y%m%d_%H%M%S.%3N")
echo "$TIMESTAMP: Conversion completed" >> $LOG

Training Output

At the end of the training process, Ultralytics creates a runs/ directory containing the training artifacts.

Typical outputs include:

• Training logs
• Validation metrics
• Loss curves
• Confusion matrix
• Model weights

But most importantly, the framework typically saves two checkpoint files:

best.pt
last.pt

They serve different purposes.

best.pt is the model you will usually want to use.

During training, after each epoch, the model is evaluated on the validation dataset. Ultralytics monitors a performance metric (by default, a fitness score derived from metrics such as mAP).

Whenever the model achieves a better validation score than any previous epoch, it overwrites best.pt.

Instead, last.pt always contains the model after the final training epoch, regardless of whether it achieved the best validation results.