Apprentissage par transfert avec couches gelées dans YOLOv5

📚 Ce guide explique comment geler les couches de YOLOv5 🚀 lors de la mise en œuvre de l'apprentissage par transfert. L'apprentissage par transfert est une technique puissante d'apprentissage automatique (ML) qui te permet de réentraîner rapidement un modèle sur de nouvelles données sans avoir à réentraîner tout le réseau à partir de zéro. En gelant les poids des couches initiales et en ne mettant à jour que les paramètres des couches ultérieures, tu peux réduire considérablement les besoins en ressources de calcul et le temps d'entraînement. Cependant, cette approche peut légèrement impacter la précision finale du modèle.

Avant de commencer

Tout d'abord, clone le dépôt YOLOv5 et installe les dépendances nécessaires listées dans requirements.txt. Assure-toi de disposer d'un environnement Python>=3.8.0 avec PyTorch>=1.8 installé. Les modèles pré-entraînés et les jeux de données requis seront téléchargés automatiquement à partir de la dernière version de YOLOv5.

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

Comment fonctionne le gel des couches

Lorsque tu gèles des couches dans un réseau neuronal, tu empêches leurs paramètres (poids et biais) d'être mis à jour durant le processus d'entraînement. Dans PyTorch, cela s'obtient en réglant l'attribut requires_grad des tenseurs de la couche sur False. Par conséquent, les gradients ne sont pas calculés pour ces couches lors de la rétropropagation, ce qui économise du calcul et de la mémoire.

Voici comment YOLOv5 implémente le gel des couches dans son script d'entraînement :

# 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

Explorer l'architecture du modèle

Comprendre la structure du modèle YOLOv5 est crucial pour décider quelles couches geler. Tu peux inspecter les noms de tous les modules et leurs paramètres en utilisant l'extrait de code Python suivant :

# 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
...
"""

L'architecture YOLOv5 consiste généralement en une backbone (couches 0-9 dans les configurations standard comme YOLOv5s/m/l/x) responsable de l'extraction de caractéristiques, et une tête (les couches restantes) qui effectue la détection d'objets.

# 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

Options de gel

Tu peux contrôler quelles couches sont gelées en utilisant l'argument --freeze dans la commande d'entraînement. Cet argument spécifie l'index du premier module non gelé ; tous les modules avant cet index verront leurs poids gelés. Utilise model.model (un nn.Sequential) pour inspecter l'ordre des modules si tu as besoin de confirmer quels indices correspondent à un bloc particulier.

Geler uniquement la backbone

Pour geler toute la backbone (couches 0 à 9), ce qui est courant lorsque l'on adapte le modèle à de nouvelles classes d'objets tout en conservant les capacités générales d'extraction de caractéristiques apprises à partir d'un grand jeu de données comme COCO :

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

Cette stratégie est efficace lorsque ton jeu de données cible partage des caractéristiques visuelles de bas niveau similaires (bords, textures) avec les données d'entraînement originales (ex. : COCO) mais contient des catégories d'objets différentes.

Geler tout sauf les couches de détection finales

Pour geler la quasi-totalité du réseau, ne laissant que les couches de convolution de sortie finales (faisant partie du module Detect, généralement le dernier module, ex. : module 24 dans YOLOv5s) entraînables :

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

Cette approche est utile lorsque tu as principalement besoin d'ajuster le modèle pour un nombre différent de classes de sortie tout en gardant intacte la grande majorité des caractéristiques apprises. Elle nécessite le moins de ressources de calcul pour le fine-tuning.

Comparaison des performances

Pour illustrer les effets du gel des couches, nous avons entraîné YOLOv5m sur le jeu de données Pascal VOC pendant 50 époques, en partant des poids pré-entraînés officiels de COCO (yolov5m.pt). Nous avons comparé trois scénarios : entraînement de toutes les couches (--freeze 0), gel de la backbone (--freeze 10), et gel de tout sauf les couches de détection 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 hyp.finetune.yaml --freeze 10

Résultats de précision

Les résultats montrent que le gel des couches peut accélérer l'entraînement de manière significative, mais peut conduire à une légère réduction de la mAP (précision moyenne) finale. L'entraînement de toutes les couches produit généralement la meilleure précision, tandis que le gel de plus de couches offre un entraînement plus rapide au prix d'une performance potentiellement inférieure.

Résultats de mAP50 d'entraînement comparant différentes stratégies de gel Comparaison de la mAP50 pendant l'entraînement

Résultats de mAP50-95 d'entraînement comparant différentes stratégies de gel Comparaison de la mAP50-95 pendant l'entraînement

YOLOv5 frozen layer training performance

*Summary table of performance metrics*

Utilisation des ressources

Geler plus de couches réduit considérablement les besoins en mémoire GPU et l'utilisation globale. Cela fait de l'apprentissage par transfert avec des couches gelées une option intéressante lorsque tu travailles avec des ressources matérielles limitées, permettant d'entraîner des modèles plus grands ou d'utiliser des tailles d'image plus grandes que ce qui serait autrement possible.

Pourcentage de mémoire GPU allouée pendant l'entraînement Mémoire GPU allouée (%)

Pourcentage d'utilisation de la mémoire GPU pendant l'entraînement Utilisation de la mémoire GPU (%)

Quand utiliser le gel des couches

Le gel des couches pendant l'apprentissage par transfert est particulièrement avantageux dans plusieurs situations :

Ressources de calcul limitées : Si tu as des contraintes sur la mémoire GPU ou la puissance de traitement.
Petits jeux de données : Lorsque ton jeu de données cible est nettement plus petit que le jeu de données de pré-entraînement original, le gel aide à prévenir le surapprentissage.
Prototypage rapide : Lorsque tu as besoin d'adapter rapidement un modèle existant à une nouvelle tâche ou à un nouveau domaine pour une évaluation initiale.
Domaines de caractéristiques similaires : Si les caractéristiques de bas niveau de ton nouveau jeu de données sont très similaires à celles du jeu de données sur lequel le modèle a été pré-entraîné.

Explore davantage les nuances de l'apprentissage par transfert dans notre entrée de glossaire et envisage des techniques comme le réglage des hyperparamètres pour optimiser les performances.

Environnements pris en charge

Ultralytics propose divers environnements prêts à l'emploi avec des dépendances essentielles comme CUDA, CuDNN, Python et PyTorch préinstallées.

Notebooks GPU gratuits :
Google Cloud : Guide de démarrage rapide GCP
Amazon : Guide de démarrage rapide AWS
Azure : Guide de démarrage rapide AzureML
Docker : Guide de démarrage rapide Docker

État du projet

Ce badge confirme que tous les tests d'Intégration Continue (CI) des GitHub Actions de YOLOv5 passent avec succès. Ces tests CI évaluent rigoureusement la fonctionnalité et la performance de YOLOv5 à travers des opérations clés : entraînement, validation, inférence, exportation et benchmarks. Ils assurent une opération cohérente et fiable sur macOS, Windows et Ubuntu, en s'exécutant automatiquement toutes les 24 heures et à chaque nouveau commit de code.

Contributors

GLglenn-jocher¹⁶ PDpderrenger¹ RIRizwanMunawar¹

Created 12 nov. 2023Updated la semaine dernière