Трансферное обучение с замороженными слоями в YOLOv5#

📚 В этом руководстве объясняется, как заморозить слои 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 обычно состоит из бэкбона (слои 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), чтобы просмотреть порядок модулей, если нужно проверить, каким индексам соответствует определённый блок.

Заморозка только бэкбона#

Чтобы заморозить весь бэкбон (слои с 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), заморозку бэкбона (--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 data/hyps/hyp.VOC.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 (%)

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

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

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

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

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

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

Бесплатные блокноты с GPU:
Google Cloud: краткое руководство по GCP
Amazon: краткое руководство по AWS
Azure: краткое руководство по AzureML
Docker: Краткое руководство по Docker

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

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

Участники

RAraimbekovm² GLglenn-jocher² PDpderrenger¹

Создано 2 месяца назадОбновлено на прошлой неделе