YOLO Training Pipeline: Difference between revisions

From LogicalDOC Community Wiki
Jump to navigationJump to search
Giuseppe (talk | contribs)
Giuseppe (talk | contribs)
Line 119: Line 119:
Example:
Example:


<pre>
<syntaxhighlight lang="yaml">
path: dataset
path: dataset
train: images/train
train: images/train
Line 129: Line 129:
   2: seller_name
   2: seller_name
   3: total
   3: total
</pre>
</syntaxhighlight>


The class definitions must match those used during annotation.
The class definitions must match those used during annotation.


=== Create a Python Training Script ===
=== Create a Python Training Script ===

Revision as of 12:54, 24 June 2026

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

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

The class definitions must match those used during annotation.

Create a Python Training Script

A training script is used to configure and start the training process.

The example script shown in this guide explicitly defines several training parameters for clarity. However, most of these parameters are optional because Ultralytics provides sensible default values.

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,  # number of epochs without sensible improvements stops the training
        imgsz=846,
        name="target",
        device=0,
        amp=True,
        save=True,
        save_period=10,
        cache="disk"
        # augment=False,
        # close_mosaic=20
    )

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

Start the Training

Create a Python training script using 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.

Execute the training script.


During training, Ultralytics YOLO:

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

The number of iterations depends on:

  • Dataset size
  • Number of epochs
  • Batch size

The batch size can be specified through the training parameters. Using -1 allows Ultralytics to select the default value automatically.

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

The most important file is:

best.pt

This file contains the trained model and should be used for subsequent validation, conversion, or deployment.