Entrena YOLOv5 con datos personalizados#

📚 Esta guía explica cómo entrenar tu propio conjunto de datos personalizado con el modelo YOLOv5 🚀. Entrenar modelos personalizados es un paso fundamental para adaptar las soluciones de visión artificial a aplicaciones concretas del mundo real, más allá de la detección de objetos genérica.

Antes de empezar#

Primero, asegúrate de que tienes configurado el entorno necesario. Clona el repositorio de YOLOv5 e instala las dependencias necesarias desde requirements.txt. Es imprescindible contar con un entorno de Python>=3.8.0 y PyTorch>=1.8. Los modelos y conjuntos de datos se descargan automáticamente de la versión más reciente de YOLOv5 si no se encuentran localmente.

git clone https://github.com/ultralytics/yolov5 # Clone the repository
cd yolov5
pip install -r requirements.txt # Install dependencies

Entrena con datos personalizados#

Desarrollar un modelo personalizado de detección de objetos es un proceso iterativo:

1. Recopila y organiza imágenes: reúne imágenes relevantes para tu tarea específica. Los datos diversos y de alta calidad son fundamentales. Consulta nuestra guía sobre recopilación y anotación de datos.
2. Etiqueta los objetos: anota con precisión los objetos de interés en tus imágenes.
3. Entrena un modelo: usa los datos etiquetados para entrenar tu modelo YOLOv5. Aprovecha el aprendizaje por transferencia partiendo de pesos preentrenados.
4. Despliega y predice: utiliza el modelo entrenado para realizar inferencias con datos nuevos que no haya visto.
5. Recopila casos límite: identifica situaciones en las que el modelo tenga un rendimiento deficiente (casos límite) y añade datos similares a tu conjunto de datos para mejorar su robustez. Repite el ciclo.

La plataforma de Ultralytics ofrece una solución simplificada sin código para todo este ciclo de operaciones de aprendizaje automático (MLOps), incluida la gestión de conjuntos de datos, el entrenamiento de modelos y el despliegue.

Antes de iniciar el entrenamiento, es esencial preparar el conjunto de datos.

1. Crea un conjunto de datos#

Los modelos YOLOv5 necesitan datos etiquetados para aprender las características visuales de las clases de objetos. Organizar correctamente tu conjunto de datos es fundamental.

1.1 Crea dataset.yaml#

El archivo de configuración del conjunto de datos (por ejemplo, coco128.yaml) describe la estructura del conjunto de datos, los nombres de las clases y las rutas a los directorios de imágenes. COCO128 es un pequeño conjunto de datos de ejemplo que contiene las primeras 128 imágenes del amplio conjunto de datos COCO. Es útil para probar rápidamente el flujo de entrenamiento y diagnosticar posibles problemas como el sobreajuste.

La estructura del archivo dataset.yaml incluye:

• path: el directorio raíz que contiene el conjunto de datos.
• train, val, test: rutas relativas desde path a directorios que contienen imágenes o archivos de texto con listas de rutas de imágenes para los conjuntos de entrenamiento, validación y prueba.
• names: un diccionario que asigna los índices de clase (a partir de 0) a los nombres de clase correspondientes.

Puedes establecer path como un directorio absoluto (por ejemplo, /home/user/datasets/coco128) o como una ruta relativa, como ../datasets/coco128, al iniciar el entrenamiento desde la raíz del repositorio de YOLOv5.

A continuación se muestra la estructura de coco128.yaml (ver en GitHub):

# Dataset root directory relative to the yolov5 directory
path: ../datasets/coco128

# Train/val/test sets: specify directories, *.txt files, or lists
train: images/train2017 # 128 images for training
val: images/train2017 # 128 images for validation
test: # Optional path to test images

# Classes (example using 80 COCO classes)
names:
    0: person
    1: bicycle
    2: car
    # ... (remaining COCO classes)
    77: teddy bear
    78: hair drier
    79: toothbrush

1.2 Aprovecha los modelos para el etiquetado automatizado#

El etiquetado manual es el método más habitual, pero lleva mucho tiempo. Los modelos fundacionales pueden automatizar o semiautomatizar la anotación y acelerar la creación de conjuntos de datos. Estos son algunos modelos que pueden ayudar a generar etiquetas:

• Google Gemini: los grandes modelos multimodales como Gemini tienen una potente capacidad para comprender imágenes. Puedes pedirles que identifiquen y localicen objetos en las imágenes, generando cuadros delimitadores o descripciones que se pueden convertir en etiquetas en formato YOLO. Descubre su potencial en el cuaderno de tutorial incluido.
• SAM2 (modelo Segment Anything 2): los modelos fundacionales centrados en la segmentación, como SAM2, pueden identificar y delimitar objetos con gran precisión. Aunque se usan principalmente para la segmentación, las máscaras resultantes suelen poder convertirse en anotaciones de cuadros delimitadores adecuadas para tareas de detección de objetos.
• YOLOWorld: este modelo ofrece capacidades de detección con vocabulario abierto. Puedes proporcionar descripciones de texto de los objetos que te interesan y YOLOWorld los localizará en las imágenes sin haber recibido entrenamiento previo con esas clases específicas. Puedes usarlo como punto de partida para generar etiquetas iniciales y luego perfeccionarlas.

El uso de estos modelos puede servir como paso de «preetiquetado» y reducir el trabajo manual necesario. Sin embargo, es fundamental revisar y perfeccionar las etiquetas generadas automáticamente para garantizar su precisión y coherencia, ya que su calidad afecta directamente al rendimiento del modelo YOLOv5 entrenado. Después de generar las etiquetas (y, si hace falta, perfeccionarlas), asegúrate de que cumplen el formato YOLO: un archivo *.txt por imagen, con cada línea representando un objeto como class_index x_center y_center width height (coordenadas normalizadas y clase indexada desde cero). Si una imagen no contiene objetos de interés, no hace falta crear el archivo *.txt correspondiente.

Las especificaciones de los archivos *.txt del formato YOLO son precisas:

• Una fila por cada cuadro delimitador.
• Cada fila debe contener: class_index x_center y_center width height.
• Las coordenadas deben estar normalizadas en un rango de 0 a 1. Para ello, divide los valores en píxeles de x_center y width entre el ancho total de la imagen, y divide y_center y height entre la altura total de la imagen.
• Los índices de clase empiezan en cero (es decir, la primera clase se representa con 0, la segunda con 1, y así sucesivamente).

El archivo de etiquetas correspondiente a la imagen anterior, con dos objetos «person» (índice de clase 0) y un objeto «tie» (índice de clase 27), tendría este aspecto:

1.3 Organiza los directorios#

Estructura el directorio de tus conjuntos de datos como se muestra a continuación. De forma predeterminada, YOLOv5 espera que el directorio del conjunto de datos (por ejemplo, /coco128) esté dentro de una carpeta /datasets situada junto al directorio del repositorio /yolov5.

YOLOv5 encuentra automáticamente las etiquetas de cada imagen sustituyendo la última aparición de /images/ en la ruta de la imagen por /labels/. Por ejemplo:

../datasets/coco128/images/im0.jpg # Path to the image file
../datasets/coco128/labels/im0.txt # Path to the corresponding label file

La estructura de directorios recomendada es:

/datasets/
└── coco128/  # Dataset root
    ├── images/
    │   ├── train2017/  # Training images
    │   │   ├── 000000000009.jpg
    │   │   └── ...
    │   └── val2017/    # Validation images (optional if using same set for train/val)
    │       └── ...
    └── labels/
        ├── train2017/  # Training labels
        │   ├── 000000000009.txt
        │   └── ...
        └── val2017/    # Validation labels (optional if using same set for train/val)
            └── ...

2. Selecciona un modelo#

Elige un modelo preentrenado para iniciar el proceso de entrenamiento. Empezar con pesos preentrenados acelera considerablemente el aprendizaje y mejora el rendimiento frente a entrenar desde cero. YOLOv5 ofrece modelos de distintos tamaños, cada uno con un equilibrio diferente entre velocidad y precisión. Por ejemplo, YOLOv5s es el segundo modelo más pequeño y el segundo más rápido, por lo que es adecuado para entornos con recursos limitados. Consulta la tabla del README para comparar en detalle todos los modelos disponibles.

3. Entrena#

Inicia el entrenamiento del modelo con el script train.py. Estos son algunos argumentos esenciales:

• --img: define el tamaño de las imágenes de entrada (por ejemplo, --img 640). Los tamaños más grandes suelen ofrecer mayor precisión, pero requieren más memoria de GPU.
• --batch: determina el tamaño del lote (por ejemplo, --batch 16). Elige el tamaño máximo que admita tu GPU.
• --epochs: especifica el número total de épocas de entrenamiento (por ejemplo, --epochs 100). Una época representa un ciclo completo por todo el conjunto de datos de entrenamiento.
• --data: ruta al archivo dataset.yaml (por ejemplo, --data coco128.yaml).
• --weights: ruta al archivo de pesos iniciales. Se recomienda encarecidamente usar pesos preentrenados (por ejemplo, --weights yolov5s.pt) para converger más rápido y obtener mejores resultados. Para entrenar desde cero (no se recomienda, salvo que tengas un conjunto de datos muy grande y necesidades específicas), usa --weights '' --cfg yolov5s.yaml.

Los pesos preentrenados se descargan automáticamente de la versión más reciente de YOLOv5 si no se encuentran localmente.

# Example: Train YOLOv5s on the COCO128 dataset for 3 epochs
python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt

Todos los resultados del entrenamiento, incluidos los pesos y los registros, se guardan en el directorio runs/train/. Cada sesión de entrenamiento crea un subdirectorio nuevo (por ejemplo, runs/train/exp, runs/train/exp2, etc.). Para una experiencia interactiva y práctica, consulta la sección de entrenamiento de nuestros cuadernos de tutorial oficiales:

4. Visualiza#

YOLOv5 se integra a la perfección con varias herramientas para visualizar el progreso del entrenamiento, evaluar los resultados y supervisar el rendimiento en tiempo real.

Registro y visualización con Comet#

Comet está totalmente integrado para realizar un seguimiento completo de los experimentos. Visualiza métricas en directo, guarda hiperparámetros, gestiona conjuntos de datos y puntos de control de modelos, y analiza predicciones del modelo con paneles personalizados de Comet interactivos.

Empezar es muy sencillo:

pip install comet_ml                                                          # 1. Install Comet library
export COMET_API_KEY=YOUR_API_KEY_HERE                                        # 2. Set your Comet API key (create a free account at Comet.ml)
python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt # 3. Train your model - Comet automatically logs everything!

Descubre en detalle las funciones compatibles en nuestra guía de integración de Comet. Consulta la documentación oficial para obtener más información sobre las capacidades de Comet. Prueba el cuaderno de Colab de Comet con esta demostración en directo:

Registro y automatización con ClearML#

La integración con ClearML permite realizar un seguimiento detallado de experimentos, gestionar las versiones de los conjuntos de datos e incluso ejecutar tareas de entrenamiento de forma remota. Activa ClearML con estos sencillos pasos:

• Instala el paquete: pip install clearml
• Inicializa ClearML: ejecuta clearml-init una vez para conectarte a tu servidor de ClearML (autoalojado o al plan gratuito).

ClearML captura automáticamente los detalles de los experimentos, las cargas de modelos, las comparaciones, los cambios de código sin confirmar y los paquetes instalados, lo que garantiza una reproducibilidad completa. Puedes programar fácilmente tareas de entrenamiento en agentes remotos y gestionar las versiones de los conjuntos de datos con ClearML Data. Consulta la guía de integración de ClearML para obtener todos los detalles.

Registro local#

Los resultados del entrenamiento se registran automáticamente con TensorBoard y se guardan como archivos CSV en el directorio del experimento correspondiente (por ejemplo, runs/train/exp). Los datos registrados incluyen:

• Pérdida de entrenamiento y validación, y métricas de rendimiento.
• Imágenes de muestra que muestran las aumentaciones aplicadas (como mosaicos).
• Etiquetas de referencia junto con predicciones del modelo para su inspección visual.
• Métricas de evaluación clave, como las curvas de precisión-recuperación (PR).
• Matrices de confusión para analizar en detalle el rendimiento por clase.

El archivo results.csv se actualiza después de cada época y se representa como results.png al finalizar el entrenamiento. También puedes representar manualmente cualquier archivo results.csv con la función de utilidad proporcionada:

from utils.plots import plot_results

# Representa los resultados de un directorio de ejecución de entrenamiento específico
plot_results("runs/train/exp/results.csv")  # Esto generará 'results.png' en el mismo directorio

5. Próximos pasos#

Al completar el entrenamiento correctamente, se guarda el punto de control del modelo con mejor rendimiento (best.pt), listo para el despliegue o para seguir perfeccionándolo. Estos son algunos de los pasos siguientes que puedes dar:

• Ejecuta inferencias con imágenes o vídeos nuevos usando el modelo entrenado mediante la CLI o Python.
• Ejecuta val.py para evaluar la precisión del modelo y su capacidad de generalización con distintas particiones de datos (por ejemplo, un conjunto de prueba reservado).
• Exporta el modelo a distintos formatos de despliegue, como ONNX, TensorFlow SavedModel o TensorRT, para optimizar las inferencias en diversas plataformas.
• Usa la evolución de hiperparámetros para intentar mejorar aún más el rendimiento.
• Sigue mejorando tu modelo con nuestros consejos para obtener los mejores resultados de entrenamiento y añade progresivamente datos más diversos y difíciles basándote en el análisis del rendimiento.

Entornos compatibles#

Ultralytics ofrece entornos listos para usar con dependencias esenciales como CUDA, cuDNN, Python y PyTorch, para que puedas empezar sin complicaciones.

• Cuadernos gratuitos con GPU:
  • 
  • 
  • 
• Plataformas en la nube:
  • Google Cloud: Guía de inicio rápido de GCP
  • Amazon AWS: guía de inicio rápido de AWS
  • Microsoft Azure: guía de inicio rápido de AzureML
• Configuración local:
  • Docker: Guía de inicio rápido de Docker

Estado del proyecto#

Esta insignia indica que todas las pruebas de integración continua (CI) de GitHub Actions de YOLOv5 se superan correctamente. Estas rigurosas pruebas de CI cubren las funciones principales, como el entrenamiento, la validación, la inferencia, la exportación y las pruebas de rendimiento, en macOS, Windows y Ubuntu. Las pruebas se ejecutan automáticamente cada 24 horas y con cada confirmación de código para garantizar una estabilidad constante y un rendimiento óptimo.

Preguntas frecuentes#