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

Expected Training Directory Structure

Before starting the training process, the training directory should be organized as follows:

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

Where:

• dataset/images/train contains the images used for training.
• dataset/images/val contains the images used for validation.
• dataset/labels/train contains the annotation files corresponding to the training images.
• dataset/labels/val contains the annotation files corresponding to the validation images.
• data.yaml defines the dataset configuration used by Ultralytics YOLO.
• synset.txt contains the list of class names used during annotation.
• notes.json is generated by Label Studio and contains additional project metadata.
• train.py defines the YOLO training configuration and launches the training process.
• train.sh (optional) automates dataset preparation, training, and model conversion.
• convert_pt_to_onnx.py converts the trained PyTorch model (`best.pt`) into the ONNX format required by LogicalDOC.

The following sections describe how to create this directory structure starting from the dataset exported by Label Studio.

Populate the Training Directory

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

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 start the training and automatically distribute the images and annotation files between the train/ and val/ directories, you can use the provided train.sh script.

Alternatively, you can manually organize the dataset by moving the desired number of images and their corresponding annotation files into the val/ directories, and then execute the train.py script directly.

To use the automated workflow, the train.sh script must be placed in the training directory. The script also invokes convert_pt_to_onnx.py to convert the trained PyTorch model (.pt) into the ONNX format required by LogicalDOC. Therefore, the `convert_pt_to_onnx.py` script must also be present in the same directory.

An example implementation of convert_pt_to_onnx.py is available in the YOLO to ONNX Conversion page.

dataset/
├── images/
│   ├── train/
│   └── val/
├── labels/
│   ├── train/
│   └── val/
├── notes.json
├── data.yaml
├── synset.txt
├── train.py
├── train.sh
└── convert-pt-to-onnx.py

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.