Transferencia de aprendizaje con capas congeladas en YOLOv5

📚 Esta guía explica cómo congelar capas de YOLOv5 🚀 al implementar transferencia de aprendizaje. La transferencia de aprendizaje es una técnica potente de aprendizaje automático (ML) que te permite reentrenar rápidamente un modelo con nuevos datos sin volver a entrenar 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 datasets 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 estableciendo el atributo requires_grad de los tensores de la capa en False. Como consecuencia, no se calculan gradientes para estas capas durante la retropropagación, ahorrando cómputo 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 usando el siguiente fragmento de 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 consiste típicamente en un backbone (capas 0-9 en configuraciones estándar como YOLOv5s/m/l/x) responsable de la extracción de características, y un head (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 usando 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 tendrá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 manteniendo las capacidades generales de extracción de características aprendidas de un dataset grande como COCO:

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

Esta estrategia es efectiva cuando tu dataset 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 finales de detección

Para congelar casi toda la red, dejando 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) entrenables:

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 ajuste fino.

Comparación de rendimiento

Para ilustrar los efectos de congelar capas, entrenamos YOLOv5m en el dataset Pascal VOC durante 50 épocas, comenzando desde los pesos oficiales preentrenados en COCO (yolov5m.pt). Comparamos tres escenarios: entrenar todas las capas (--freeze 0), congelar el backbone (--freeze 10) y congelar todo excepto las capas de detección final (--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 hyp.finetune.yaml --freeze 10

Resultados de precisión

Los resultados muestran que congelar capas puede acelerar significativamente el entrenamiento pero puede llevar a una ligera reducción en la mAP (precisión media media) 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 la transferencia de aprendizaje con capas congeladas sea una opción atractiva cuando trabajas con recursos de hardware limitados, permitiendo entrenar modelos más grandes o usar tamaños de imagen mayores de lo que sería posible de otra manera.

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

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

Cuándo usar 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.
Datasets pequeños: Cuando tu dataset objetivo es significativamente más pequeño que el dataset original de preentrenamiento, congelar ayuda a prevenir el sobreajuste.
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 dataset son muy similares a las del dataset en el que el modelo fue preentrenado.

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

Entornos compatibles

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: Guía de inicio rápido de AWS
Azure: Guía de inicio rápido de AzureML
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 están pasando 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. Aseguran un funcionamiento consistente y confiable en macOS, Windows y Ubuntu, ejecutándose automáticamente cada 24 horas y con cada nuevo commit de código.

Contributors

GLglenn-jocher¹⁶ PDpderrenger¹ RIRizwanMunawar¹

Created 12 nov 2023Updated la semana pasada