AI Model Training with Ultralytics YOLO
Introduction
Training a deep learning model involves feeding it data and adjusting its parameters so that it can make accurate predictions. Train mode in Ultralytics YOLO11 is engineered for effective and efficient training of object detection models, fully utilizing modern hardware capabilities. This guide aims to cover all the details you need to get started with training your own models using YOLO11's robust set of features.
Watch: How to Train a YOLO model on Your Custom Dataset in Google Colab.
Why Choose Ultralytics YOLO for Training?
Here are some compelling reasons to opt for YOLO11's Train mode:
- Efficiency: Make the most out of your hardware, whether you're on a single-GPU setup or scaling across multiple GPUs.
- Versatility: Train on custom datasets in addition to readily available ones like COCO, VOC, and ImageNet.
- User-Friendly: Simple yet powerful CLI and Python interfaces for a straightforward training experience.
- Hyperparameter Flexibility: A broad range of customizable hyperparameters to fine-tune model performance.
Key Features of Train Mode
The following are some notable features of YOLO11's Train mode:
- Automatic Dataset Download: Standard datasets like COCO, VOC, and ImageNet are downloaded automatically on first use.
- Multi-GPU Support: Scale your training efforts seamlessly across multiple GPUs to expedite the process.
- Hyperparameter Configuration: The option to modify hyperparameters through YAML configuration files or CLI arguments.
- Visualization and Monitoring: Real-time tracking of training metrics and visualization of the learning process for better insights.
Tip
- YOLO11 datasets like COCO, VOC, ImageNet and many others automatically download on first use, i.e.
yolo train data=coco.yaml
Usage Examples
Train YOLO11n on the COCO8 dataset for 100 epochs at image size 640. The training device can be specified using the device
argument. If no argument is passed GPU device=0
will be used if available, otherwise device='cpu'
will be used. See Arguments section below for a full list of training arguments.
Windows Multi-Processing Error
On Windows, you may receive a RuntimeError
when launching the training as a script. Add a if __name__ == "__main__":
block before your training code to resolve it.
Single-GPU and CPU Training Example
Device is determined automatically. If a GPU is available then it will be used, otherwise training will start on CPU.
from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.yaml") # build a new model from YAML
model = YOLO("yolo11n.pt") # load a pretrained model (recommended for training)
model = YOLO("yolo11n.yaml").load("yolo11n.pt") # build from YAML and transfer weights
# Train the model
results = model.train(data="coco8.yaml", epochs=100, imgsz=640)
Multi-GPU Training
Multi-GPU training allows for more efficient utilization of available hardware resources by distributing the training load across multiple GPUs. This feature is available through both the Python API and the command-line interface. To enable multi-GPU training, specify the GPU device IDs you wish to use.
Multi-GPU Training Example
To train with 2 GPUs, CUDA devices 0 and 1 use the following commands. Expand to additional GPUs as required.
from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt") # load a pretrained model (recommended for training)
# Train the model with 2 GPUs
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])
Apple Silicon MPS Training
With the support for Apple silicon chips integrated in the Ultralytics YOLO models, it's now possible to train your models on devices utilizing the powerful Metal Performance Shaders (MPS) framework. The MPS offers a high-performance way of executing computation and image processing tasks on Apple's custom silicon.
To enable training on Apple silicon chips, you should specify 'mps' as your device when initiating the training process. Below is an example of how you could do this in Python and via the command line:
MPS Training Example
from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt") # load a pretrained model (recommended for training)
# Train the model with MPS
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device="mps")
While leveraging the computational power of the Apple silicon chips, this enables more efficient processing of the training tasks. For more detailed guidance and advanced configuration options, please refer to the PyTorch MPS documentation.
Resuming Interrupted Trainings
Resuming training from a previously saved state is a crucial feature when working with deep learning models. This can come in handy in various scenarios, like when the training process has been unexpectedly interrupted, or when you wish to continue training a model with new data or for more epochs.
When training is resumed, Ultralytics YOLO loads the weights from the last saved model and also restores the optimizer state, learning rate scheduler, and the epoch number. This allows you to continue the training process seamlessly from where it was left off.
You can easily resume training in Ultralytics YOLO by setting the resume
argument to True
when calling the train
method, and specifying the path to the .pt
file containing the partially trained model weights.
Below is an example of how to resume an interrupted training using Python and via the command line:
Resume