Aprendizaje por transferencia con capas congeladas en YOLOv5#

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

Antes de empezar#

Primero, clona el repositorio de YOLOv5 e instala las dependencias necesarias que se indican en requirements.txt. Asegúrate de tener un entorno de 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 versión más reciente 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, impides que sus parámetros (pesos y sesgos) se actualicen durante el proceso de entrenamiento. En PyTorch, esto se consigue 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, lo que ahorra capacidad de cálculo y memoria.

Así implementa YOLOv5 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 fundamental para decidir qué capas congelar. Puedes inspeccionar los nombres de todos los módulos y sus parámetros con 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 normalmente consta 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 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 mediante 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 concreto.

Congelar solo el backbone#

Para congelar todo el backbone (capas 0 a 9), algo habitual al adaptar el modelo a nuevas clases de objetos mientras se conservan las capacidades generales de extracción de características aprendidas a partir de un conjunto de datos grande como COCO:

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

Esta estrategia resulta eficaz cuando el conjunto de datos de destino comparte características visuales de bajo nivel similares (bordes y 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 prácticamente toda la red y dejar entrenables únicamente las capas convolucionales de salida finales (que forman 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 resulta útil cuando principalmente necesitas ajustar el modelo a 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 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 las capas salvo las capas finales de detección (--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 considerablemente el entrenamiento, pero puede provocar una ligera reducción de la mAP (precisión media promedio) final. Entrenar todas las capas suele ofrecer la mejor precisión, mientras que congelar más capas permite entrenar más rápido a costa de un rendimiento potencialmente menor.

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

Resultados de mAP50-95 del entrenamiento comparando distintas 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 convierte el aprendizaje por transferencia con capas congeladas en una opción atractiva cuando trabajas con recursos de hardware limitados, ya que permite entrenar modelos más grandes o utilizar tamaños de imagen mayores de lo que sería posible en otras circunstancias.

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

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

Cuándo usar la congelación de capas#

La congelación de capas durante el aprendizaje por transferencia resulta especialmente ventajosa en varias situaciones:

Recursos computacionales limitados: si tienes limitaciones de memoria de GPU o potencia de procesamiento.
Conjuntos de datos pequeños: cuando tu conjunto de datos de destino es considerablemente más pequeño que el conjunto de datos original de preentrenamiento, la congelación ayuda a evitar el sobreajuste.
Prototipado rápido: cuando necesitas adaptar rápidamente un modelo existente a una tarea o dominio nuevos para realizar una evaluación inicial.
Dominios de características similares: si las características de bajo nivel de tu nuevo conjunto de datos son muy similares a las del conjunto de datos con el que se preentrenó el modelo.

Descubre más sobre los matices del aprendizaje por transferencia en nuestra entrada del glosario y considera técnicas como el 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.

Cuadernos de GPU gratuitos:
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 GitHub Actions de YOLOv5 se superan correctamente. 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 uniforme y fiable en macOS, Windows y Ubuntu, y se ejecutan automáticamente cada 24 horas y con cada nuevo commit de código.

Colaboradores

RAraimbekovm² GLglenn-jocher² PDpderrenger¹

Creado hace 2 mesesActualizado la semana pasada