Transferencia de aprendizaje con capas congeladas en YOLOv5#

📚 Esta guía explica cómo congelar las capas de YOLOv5 🚀 al implementar transfer learning. El transfer learning es una potente técnica de aprendizaje automático (ML) que te permite reentrenar rápidamente un modelo con nuevos datos sin tener que reentrenar toda la red desde cero. Al congelar los pesos de las capas iniciales y actualizar solo los parámetros de las capas posteriores, puedes reducir significativamente los requisitos de recursos computacionales y el tiempo de entrenamiento. Sin embargo, este enfoque podría afectar ligeramente la precisión final del modelo.

Antes de empezar#

Primero, clona el repositorio de YOLOv5 e instala las dependencias necesarias enumeradas en requirements.txt. Asegúrate de tener un entorno Python>=3.8.0 con PyTorch>=1.8 instalado. Los modelos preentrenados y los conjuntos de datos necesarios se descargarán automáticamente desde la última versión de YOLOv5.

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

Cómo funciona la congelación de capas#

Cuando congelas capas en una red neuronal, evitas que sus parámetros (pesos y sesgos) se actualicen durante el proceso de entrenamiento. En PyTorch, esto se logra configurando el atributo requires_grad de los tensores de la capa en False. En consecuencia, no se calculan gradientes para estas capas durante la propagación hacia atrás, lo que ahorra computación y memoria.

Así es como YOLOv5 implementa la congelación de capas en su script de entrenamiento:

# Freeze specified layers
freeze = [f"model.{x}." for x in range(freeze)]  # Define layers to freeze based on module index
for k, v in model.named_parameters():
    v.requires_grad = True  # Ensure all parameters are initially trainable
    if any(x in k for x in freeze):
        print(f"Freezing layer: {k}")
        v.requires_grad = False  # Disable gradient calculation for frozen layers

Exploración de la arquitectura del modelo#

Comprender la estructura del modelo YOLOv5 es crucial para decidir qué capas congelar. Puedes inspeccionar los nombres de todos los módulos y sus parámetros utilizando el siguiente fragmento de código Python:

# Assuming 'model' is your loaded YOLOv5 model instance
for name, param in model.named_parameters():
    print(name)

"""
Example Output:
model.0.conv.conv.weight
model.0.conv.bn.weight
model.0.conv.bn.bias
model.1.conv.weight
model.1.bn.weight
model.1.bn.bias
model.2.cv1.conv.weight
model.2.cv1.bn.weight
...
"""

La arquitectura de YOLOv5 consta típicamente de un backbone (capas 0-9 en configuraciones estándar como YOLOv5s/m/l/x) responsable de la extracción de características, y una cabecera (las capas restantes) que realiza la detección de objetos.

# Example YOLOv5 v6.0 backbone structure
backbone:
    # [from, number, module, args]
    - [-1, 1, Conv, [64, 6, 2, 2]]  # Layer 0: Initial convolution (P1/2 stride)
    - [-1, 1, Conv, [128, 3, 2]] # Layer 1: Downsampling convolution (P2/4 stride)
    - [-1, 3, C3, [128]]          # Layer 2: C3 module
    - [-1, 1, Conv, [256, 3, 2]] # Layer 3: Downsampling convolution (P3/8 stride)
    - [-1, 6, C3, [256]]          # Layer 4: C3 module
    - [-1, 1, Conv, [512, 3, 2]] # Layer 5: Downsampling convolution (P4/16 stride)
    - [-1, 9, C3, [512]]          # Layer 6: C3 module
    - [-1, 1, Conv, [1024, 3, 2]]# Layer 7: Downsampling convolution (P5/32 stride)
    - [-1, 3, C3, [1024]]         # Layer 8: C3 module
    - [-1, 1, SPPF, [1024, 5]]    # Layer 9: Spatial Pyramid Pooling Fast

# Example YOLOv5 v6.0 head structure
head:
    - [-1, 1, Conv, [512, 1, 1]] # Layer 10
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]] # Layer 11
    - [[-1, 6], 1, Concat, [1]] # Layer 12: Concatenate with backbone P4 (from layer 6)
    - [-1, 3, C3, [512, False]] # Layer 13: C3 module
    # ... subsequent head layers for feature fusion and detection

Opciones de congelación#

Puedes controlar qué capas se congelan utilizando el argumento --freeze en el comando de entrenamiento. Este argumento especifica el índice del primer módulo no congelado; todos los módulos anteriores a este índice verán sus pesos congelados. Usa model.model (un nn.Sequential) para inspeccionar el orden de los módulos si necesitas confirmar qué índices corresponden a un bloque en particular.

Congelar solo el backbone#

Para congelar todo el backbone (capas 0 a 9), lo cual es común al adaptar el modelo a nuevas clases de objetos mientras se conservan las capacidades generales de extracción de características aprendidas de un gran conjunto de datos como COCO:

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 10

Esta estrategia es eficaz cuando tu conjunto de datos objetivo comparte características visuales de bajo nivel similares (bordes, texturas) con los datos de entrenamiento originales (por ejemplo, COCO), pero contiene categorías de objetos diferentes.

Congelar todo excepto las capas de detección final#

Para congelar casi toda la red, dejando entrenables solo las capas de convolución de salida final (parte del módulo Detect, normalmente el último módulo, por ejemplo, el módulo 24 en YOLOv5s):

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 24

Este enfoque es útil cuando necesitas principalmente ajustar el modelo para un número diferente de clases de salida manteniendo intacta la gran mayoría de las características aprendidas. Requiere la menor cantidad de recursos computacionales para el fine-tuning.

Comparación de rendimiento#

Para ilustrar los efectos de congelar capas, entrenamos YOLOv5m en el conjunto de datos Pascal VOC durante 50 épocas, partiendo de los pesos preentrenados oficiales de COCO (yolov5m.pt). Comparamos tres escenarios: entrenar todas las capas (--freeze 0), congelar el backbone (--freeze 10) y congelar todas excepto las capas de detección finales (--freeze 24).

# Example command for training with backbone frozen
python train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp data/hyps/hyp.VOC.yaml --freeze 10

Resultados de precisión#

Los resultados muestran que congelar capas puede acelerar significativamente el entrenamiento, pero puede provocar una ligera reducción en el mAP (mean Average Precision) final. Entrenar todas las capas generalmente produce la mejor precisión, mientras que congelar más capas ofrece un entrenamiento más rápido a costa de un rendimiento potencialmente menor.

Resultados de mAP50 de entrenamiento comparando diferentes estrategias de congelación Comparación de mAP50 durante el entrenamiento

Resultados de mAP50-95 de entrenamiento comparando diferentes estrategias de congelación Comparación de mAP50-95 durante el entrenamiento

YOLOv5 frozen layer training performance

*Summary table of performance metrics*

Utilización de recursos#

Congelar más capas reduce sustancialmente los requisitos de memoria de la GPU y la utilización general. Esto hace que el transfer learning con capas congeladas sea una opción atractiva cuando se trabaja con recursos de hardware limitados, permitiendo entrenar modelos más grandes o usar tamaños de imagen mayores de lo que de otro modo sería posible.

Porcentaje de memoria de GPU asignada durante el entrenamiento Memoria de GPU asignada (%)

Porcentaje de utilización de memoria de GPU durante el entrenamiento Utilización de GPU (%)

Cuándo utilizar la congelación de capas#

La congelación de capas durante la transferencia de aprendizaje es particularmente ventajosa en varias situaciones:

Recursos computacionales limitados: Si tienes restricciones en la memoria de la GPU o en la potencia de procesamiento.
Conjuntos de datos pequeños: Cuando tu conjunto de datos de destino es significativamente más pequeño que el conjunto de datos de preentrenamiento original, la congelación ayuda a prevenir el overfitting.
Prototipado rápido: Cuando necesitas adaptar rápidamente un modelo existente a una nueva tarea o dominio para una evaluación inicial.
Dominios de características similares: Si las características de bajo nivel en tu nuevo conjunto de datos son muy similares a las del conjunto de datos en el que se preentrenó el modelo.

Explora más sobre los matices del transfer learning en nuestra entrada del glosario y considera técnicas como el ajuste de hiperparámetros para optimizar el rendimiento.

Entornos admitidos#

Ultralytics ofrece varios entornos listos para usar con dependencias esenciales como CUDA, CuDNN, Python y PyTorch preinstaladas.

Notebooks con GPU gratuita:
Google Cloud: Guía de inicio rápido de GCP
Amazon: AWS Quickstart Guide
Azure: AzureML Quickstart Guide
Docker: Guía de inicio rápido de Docker

Estado del proyecto#

Esta insignia confirma que todas las pruebas de Integración Continua (CI) de YOLOv5 GitHub Actions se están superando con éxito. Estas pruebas de CI evalúan rigurosamente la funcionalidad y el rendimiento de YOLOv5 en operaciones clave: entrenamiento, validación, inferencia, exportación y benchmarks. Garantizan un funcionamiento coherente y fiable en macOS, Windows y Ubuntu, ejecutándose automáticamente cada 24 horas y en cada nuevo commit de código.

Colaboradores

GLglenn-jocher² RAraimbekovm¹

Creado hace 2 mesesActualizado el mes pasado