Mastering Object Detection: Training YOLO on Custom Objects

By F. Shane Alvares

Now that we have our dataset ready, it’s time to train YOLO to detect custom objects. Training a deep-learning model from scratch might sound daunting, but we’ll break it down into manageable steps. By the end of this guide, you’ll know how to fine-tune YOLO, adjust hyperparameters, and optimize training for better results.

Step 1: Preparing the Dataset for Training

Before training YOLO, we need to structure our dataset properly. If you followed Part 3, you should have a well-annotated dataset from Label Studio. Now, let’s ensure it’s in the correct format:

• Convert Annotations to YOLO Format: YOLO requires annotations in a specific format: each image gets a corresponding text file where each line represents an object with its class ID and bounding box coordinates (normalized between 0 and 1).
• Organizing Files: Create three directories inside your dataset folder:

- images/train (for training images)

- images/val (for validation images)

- labels/train and labels/val (for corresponding annotation files)

• Split the Dataset: Typically, we use an 80–20 split between training and validation data. You can use Python scripts to automate this.

Step 2: Configuring YOLO for Training

Now, we need to modify some configuration files so YOLO knows about our custom objects:

• Create a Data Configuration File: This .yaml file tells YOLO where to find training and validation data. It should look something like this:

train: /path/to/dataset/images/train val: /path/to/dataset/images/val nc: <number_of_classes> names: ["class1", "class2", "class3"

• Modify the Model Configuration: Depending on the YOLO version you’re using, update the yaml or .cfg file to specify the number of classes in the last layer.
• Set Up Hyperparameters: Adjust batch size, learning rate, and epochs in the training script. A batch size of 16–32 and a learning rate of 0.01 are good starting points.

Step 3: Running the Training Process

With the configuration set, we can now start training YOLO on our custom dataset.

1. Open a terminal and navigate to the YOLO directory.
2. Run the training command (for YOLOv11, it might look something like this):

python train.py --data custom_dataset.yaml --cfg yolov11.yaml --weights yolov11.pt --epochs 100 --batch-size 16 --device 

Additional Details: Here’s a breakdown of the parameters in your YOLO training command:

python train.py --data custom_dataset.yaml --cfg yolov11.yaml --weights yolov11.pt --epochs 100 --batch-size 16 --device 0

Additional Details: Parameter Breakdown

python train.py
Runs the training script (train.py), which starts the model training process.

--data custom_dataset.yaml
Specifies the dataset configuration file. This YAML file contains paths to training/validation images, number of classes, and class names.

--cfg yolov11.yaml
Defines the model configuration file. This file contains the YOLO network architecture settings, including layers, anchors, and other hyperparameters.

--weights yolov11.pt
Specifies the initial model weights.
If using a pre-trained model (e.g., yolov11.pt), it helps with transfer learning.
If set to """ --weights '' """, the model will train from scratch.

--epochs 100
The number of training iterations over the entire dataset.
More epochs generally improve accuracy but increase training time.

--batch-size 16
The number of images processed simultaneously during training.
A higher batch size speeds up training but requires more GPU memory.

--device 0
Specifies which device to use for training:
  0: Use GPU (if available).
  cpu: Train using the CPU (much slower).
  0,1,2,...: Multi-GPU training (if multiple GPUs are available).

3. Monitor the training process. You should see metrics like loss, mAP (mean Average Precision), and training speed.

4. If loss doesn’t decrease, consider adjusting the learning rate or augmenting the dataset.

Step 4: Evaluating Training Performance

Once training completes, we need to evaluate performance:

• Check Precision and Recall: Look at the confusion matrix to see how well the model distinguishes between classes.
• Analyze mAP Score: The higher the mean Average Precision, the better the model’s accuracy.
• Test on New Images: Run inference on unseen images to visually inspect results.

Step 5: Fine-Tuning and Optimizing

Improving model performance often requires some tweaking:

• Data Augmentation: Apply techniques like flipping, rotation, and brightness adjustments.
• Hyperparameter Tuning: Experiment with different batch sizes, learning rates, and optimizer functions.
• Transfer Learning: If starting from scratch doesn’t yield good results, use pre-trained YOLO weights for better performance.

Next Steps

Now that we have a trained model, we’re ready to put it to work! In the next part of this series, we’ll explore how to evaluate and improve model accuracy further. Stay tuned!