Перейти к содержимому

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

📚 Это руководство объясняет, как заморозить YOLOv5 🚀 слои при трансферном обучении. Трансферное обучение - это полезный способ быстро переобучить модель на новых данных без необходимости переобучать всю сеть. Вместо этого часть начальных весов замораживается на месте, а остальные веса используются для вычисления потерь и обновляются оптимизатором. Это требует меньше ресурсов, чем обычное обучение, и позволяет ускорить время обучения, хотя и может привести к снижению конечной точности обучения.

Прежде чем начать

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

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

Заморозьте заднюю стенку

Все слои, которые соответствуют train.py freeze список в train.py будет заморожен путем установки их градиентов в ноль перед началом тренировки.

# Freeze
freeze = [f'model.{x}.' for x in range(freeze)]  # layers to freeze
for k, v in model.named_parameters():
    v.requires_grad = True  # train all layers
    if any(x in k for x in freeze):
        print(f'freezing {k}')
        v.requires_grad = False

Чтобы просмотреть список названий модулей:

for k, v in model.named_parameters():
    print(k)

"""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
...
model.23.m.0.cv2.bn.weight
model.23.m.0.cv2.bn.bias
model.24.m.0.weight
model.24.m.0.bias
model.24.m.1.weight
model.24.m.1.bias
model.24.m.2.weight
model.24.m.2.bias
"""

Взглянув на архитектуру модели, мы видим, что основой модели являются уровни 0-9:

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  - [-1, 1, Conv, [64, 6, 2, 2]]  # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  - [-1, 3, C3, [128]]
  - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  - [-1, 6, C3, [256]]
  - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  - [-1, 9, C3, [512]]
  - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  - [-1, 3, C3, [1024]]
  - [-1, 1, SPPF, [1024, 5]]  # 9


# YOLOv5 v6.0 head
head:
  - [-1, 1, Conv, [512, 1, 1]]
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
  - [-1, 3, C3, [512, False]]  # 13

  - [-1, 1, Conv, [256, 1, 1]]
  - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
  - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
  - [-1, 3, C3, [256, False]]  # 17 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 14], 1, Concat, [1]]  # cat head P4
  - [-1, 3, C3, [512, False]]  # 20 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]]  # cat head P5
  - [-1, 3, C3, [1024, False]]  # 23 (P5/32-large)

  - [[17, 20, 23], 1, Detect, [nc]]  # Detect(P3, P4, P5)

Так что мы можем определить, что список заморозки будет содержать все модули с 'model.0.' - 'model.9.' в своих названиях:

python train.py --freeze 10

Заморозьте все слои

Чтобы заморозить всю модель, кроме финальных выходных сверточных слоев в Detect(), мы задаем в списке freeze все модули с 'model.0.' - 'model.23.' в своих названиях:

python train.py --freeze 24

Результаты

Мы тренируем YOLOv5m на VOC на обоих вышеупомянутых сценариях, а также на модели по умолчанию (без заморозки), начиная с официальной предварительной тренировки COCO --weights yolov5m.pt:

train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp hyp.finetune.yaml

Сравнение точности

Результаты показывают, что замораживание ускоряет обучение, но немного снижает итоговую точность.

Результаты заморозки тренировки mAP50

Результаты заморозки тренировки mAP50-95

Результаты таблицы

Сравнение использования графического процессора

Интересно, что чем больше модулей заморожено, тем меньше памяти GPU требуется для обучения и тем ниже загрузка GPU. Это говорит о том, что большие модели или модели, обученные на больших -image-size, могут выиграть от замораживания, чтобы обучаться быстрее.

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

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

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

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

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

YOLOv5 CI

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



Создано 2023-11-12, Обновлено 2024-01-14
Авторы: glenn-jocher (4)

Комментарии