Прунинг и разреженность моделей в YOLOv5

📚 Это руководство объясняет, как применить прунинг к моделям 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

Проверка базовой производительности

Перед прунингом установи базовую производительность для сравнения. Эта команда тестирует YOLOv5x на COCO val2017 при размере изображения 640 пикселей. yolov5x.pt — самая большая и точная доступная модель. Другие варианты — yolov5s.pt, yolov5m.pt и yolov5l.pt, или твой собственный чекпоинт после обучения на пользовательском наборе данных ./weights/best.pt. Подробности обо всех доступных моделях смотри в таблице в README.

python val.py --weights yolov5x.pt --data coco.yaml --img 640 --half

Вывод:

val: data=/content/yolov5/data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True, dnn=False
YOLOv5 🚀 v6.0-224-g4c40933 torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)

Fusing layers...
Model Summary: 444 layers, 86705005 parameters, 0 gradients
val: Scanning '/content/datasets/coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupt: 100% 5000/5000 [00:00<?, ?it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 157/157 [01:12<00:00,  2.16it/s]
                 all       5000      36335      0.732      0.628      0.683      0.496
Speed: 0.1ms pre-process, 5.2ms inference, 1.7ms NMS per image at shape (32, 3, 640, 640)  # <--- base speed

Evaluating pycocotools mAP... saving runs/val/exp-2/yolov5x_predictions.json...
...
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.507  # <--- base mAP
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.689
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.552
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.345
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.559
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.652
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.381
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.630
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.682
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.526
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.731
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.829
Results saved to runs/val/exp

Применение прунинга к YOLOv5x (30% разреженности)

Мы можем применить прунинг к модели, используя команду torch_utils.prune(), определенную в utils/torch_utils.py. Чтобы протестировать модель после прунинга, мы обновляем val.py, чтобы провести прунинг YOLOv5x до уровня разреженности 0.3 (30% весов устанавливаются в ноль):

YOLOv5 model pruning to 30% sparsity code

Результат с 30% прунингом:

val: data=/content/yolov5/data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True, dnn=False
YOLOv5 🚀 v6.0-224-g4c40933 torch 1.10.0+cu111 CUDA:0 (Tesla V100-SXM2-16GB, 16160MiB)

Fusing layers...
Model Summary: 444 layers, 86705005 parameters, 0 gradients
Pruning model...  0.3 global sparsity
val: Scanning '/content/datasets/coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupt: 100% 5000/5000 [00:00<?, ?it/s]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 157/157 [01:11<00:00,  2.19it/s]
                 all       5000      36335      0.724      0.614      0.671      0.478
Speed: 0.1ms pre-process, 5.2ms inference, 1.7ms NMS per image at shape (32, 3, 640, 640)  # <--- prune speed

Evaluating pycocotools mAP... saving runs/val/exp-3/yolov5x_predictions.json...
...
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.489  # <--- prune mAP
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.677
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.537
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.334
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.542
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.635
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.370
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.612
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.664
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.496
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.722
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.803
Results saved to runs/val/exp-3

Анализ результатов

Из полученных результатов мы можем заметить:

  • Достигнута разреженность 30%: 30% весовых параметров модели в слоях nn.Conv2d теперь равны нулю
  • Время инференса остается неизменным: несмотря на прунинг, скорость обработки практически такая же
  • Минимальное влияние на производительность: mAP снизился незначительно с 0.507 до 0.489 (всего на 3.6%)
  • Уменьшение размера модели: модель после прунинга требует меньше памяти для хранения

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

Дообучение моделей после прунинга

Для достижения наилучших результатов модели после прунинга следует дообучить, чтобы восстановить точность. Это можно сделать путем:

  1. Применения прунинга с желаемым уровнем разреженности
  2. Обучения модели после прунинга в течение нескольких эпох с более низкой скоростью обучения (learning rate)
  3. Оценки дообученной модели после прунинга относительно базовой

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

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

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

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

YOLOv5 CI

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

Комментарии