Редактировать страницу на GitHub (EN)

Обучение методом переноса (transfer learning) с заморозкой слоев в YOLOv5

📚 Это руководство объясняет, как заморозить (freeze) слои YOLOv5 🚀 при использовании обучения методом переноса. Обучение методом переноса — это мощная техника машинного обучения (ML), позволяющая быстро переобучить модель на новых данных без обучения всей сети с нуля. Замораживая веса начальных слоев и обновляя только параметры последующих, ты можешь значительно снизить требования к вычислительным ресурсам и время обучения. Однако такой подход может незначительно повлиять на итоговую точность модели.

Перед началом

Сначала клонируй репозиторий YOLOv5 и установи необходимые зависимости, указанные в файле requirements.txt. Убедись, что у тебя установлена среда Python>=3.8.0 с PyTorch>=1.8. Предобученные модели и необходимые наборы данных будут загружены автоматически из последнего релиза YOLOv5.

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

Как работает заморозка слоев

Когда ты замораживаешь слои в нейронной сети, ты предотвращаешь обновление их параметров (весов и смещений) во время процесса обучения. В PyTorch это достигается путем установки атрибута requires_grad тензоров слоя в значение False. В результате градиенты для этих слоев не вычисляются во время обратного распространения ошибки, что экономит вычислительные ресурсы и память.

Вот как YOLOv5 реализует заморозку слоев в своем скрипте обучения:

# 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

Изучение архитектуры модели

Понимание структуры модели YOLOv5 имеет решающее значение для решения о том, какие слои замораживать. Ты можешь просмотреть названия всех модулей и их параметры, используя следующий фрагмент кода на 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
...
"""

Архитектура YOLOv5 обычно состоит из backbone (слои 0-9 в стандартных конфигурациях, таких как YOLOv5s/m/l/x), отвечающего за извлечение признаков, и головы (остальные слои), которая выполняет детекцию объектов.

# 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

Параметры заморозки

Ты можешь управлять тем, какие слои заморожены, используя аргумент --freeze в команде обучения. Этот аргумент указывает индекс первого незамороженного модуля; у всех модулей перед этим индексом веса будут заморожены. Используй model.model (объект nn.Sequential), чтобы проверить порядок модулей, если тебе нужно подтвердить, какие индексы соответствуют конкретному блоку.

Заморозка только Backbone

Чтобы заморозить весь backbone (слои с 0 по 9), что обычно делается при адаптации модели к новым классам объектов с сохранением общих возможностей извлечения признаков, изученных на большом наборе данных, таком как COCO:

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

Эта стратегия эффективна, когда твой целевой набор данных имеет схожие низкоуровневые визуальные признаки (границы, текстуры) с исходными данными для обучения (например, COCO), но содержит другие категории объектов.

Заморозка всего, кроме финальных слоев детекции

Чтобы заморозить почти всю сеть, оставив обучаемыми только финальные сверточные выходные слои (часть модуля Detect, обычно это последний модуль, например, модуль 24 в YOLOv5s):

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

Этот подход полезен, когда тебе нужно в основном настроить модель под другое количество выходных классов, сохранив при этом подавляющее большинство изученных признаков. Он требует наименьшего количества вычислительных ресурсов для тонкой настройки.

Сравнение производительности

Чтобы проиллюстрировать эффекты заморозки слоев, мы обучили YOLOv5m на наборе данных Pascal VOC в течение 50 эпох, начав с официальных предобученных весов COCO (yolov5m.pt). Мы сравнили три сценария: обучение всех слоев (--freeze 0), заморозка backbone (--freeze 10) и заморозка всего, кроме финальных слоев детекции (--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

Результаты точности

Результаты показывают, что заморозка слоев может значительно ускорить обучение, но может привести к небольшому снижению итогового mAP (средней точности). Обучение всех слоев обычно дает наилучшую точность, тогда как заморозка большего количества слоев ускоряет процесс ценой потенциально более низких показателей.

Результаты обучения mAP50 при сравнении разных стратегий заморозки Сравнение mAP50 во время обучения

Результаты обучения mAP50-95 при сравнении разных стратегий заморозки Сравнение mAP50-95 во время обучения

YOLOv5 frozen layer training performance *Summary table of performance metrics*

Использование ресурсов

Заморозка большего количества слоев существенно снижает требования к памяти GPU и общую загрузку. Это делает обучение методом переноса с заморозкой слоев привлекательным вариантом при ограниченных аппаратных ресурсах, позволяя обучать более крупные модели или использовать изображения большего размера, чем это могло бы быть возможно в противном случае.

Процент выделенной памяти GPU во время обучения Выделенная память GPU (%)

Процент использования памяти GPU во время обучения Использование GPU (%)

Когда использовать заморозку слоев

Заморозка слоев во время обучения методом переноса особенно выгодна в нескольких ситуациях:

  1. Ограниченные вычислительные ресурсы: Если у тебя есть ограничения по памяти GPU или вычислительной мощности.
  2. Маленькие наборы данных: Если твой целевой набор данных значительно меньше исходного набора данных, на котором проводилось предварительное обучение, заморозка помогает предотвратить переобучение.
  3. Быстрое прототипирование: Когда тебе нужно быстро адаптировать существующую модель к новой задаче или домену для первичной оценки.
  4. Схожие домены признаков: Если низкоуровневые признаки в твоем новом наборе данных очень похожи на те, что были в наборе данных, на котором модель была предварительно обучена.

Узнай больше о нюансах обучения методом переноса в нашей статье в глоссарии и рассмотри такие методы, как настройка гиперпараметров для оптимизации производительности.

Поддерживаемые среды

Ultralytics предлагает различные готовые к использованию среды с предустановленными необходимыми зависимостями, такими как CUDA, CuDNN, Python и PyTorch.

Статус проекта

Статус непрерывной интеграции YOLOv5

Этот бейдж подтверждает, что все тесты непрерывной интеграции (CI) GitHub Actions для YOLOv5 успешно проходят. Эти CI-тесты строго оценивают функциональность и производительность YOLOv5 по ключевым операциям: обучение, валидация, инференс, экспорт и бенчмарки. Они обеспечивают стабильную и надежную работу на macOS, Windows и Ubuntu, запускаясь автоматически каждые 24 часа и при каждом новом коммите кода.

Contributors
Редактировать страницу на GitHub (EN)
GLglenn-jocher16PDpderrenger1RIRizwanMunawar1

Комментарии