Сравнение архитектур нейросетей, основанных на детекции объектов на изображениях

Датасет был взят с Kaggle

Из-за несбалансированности классов был написан класс датасета PyTorch, который возвращает все выборки (train, valid, test) со стратифицированным разбиением по классам, а так же имеет атрибут df, который хранит в себе объект pandas.DataFrame для конкретной выборки

Faster R-CNN

Данная модель является двуступенчатой, поэтому требует много вычислительной мощности, но при этом даёт качественные результаты.

Ноутбук со всеми вычислениями и кодом тут

Использовался fine-tuning предобученной сети из PyTorch.

Обучение производилось на 3-ёх эпохах. Learning rate ($\eta$) был выбран довольно низким $2 \times 10^{-6}$

DataLoader для обучающей выборки был сэмплирован на основе веса каждого изображения. Подробнее в ноутбуке

Пример результата работы данной сети:

YOLOv8

Данная модель является более быстрой, чем Faster R-CNN. Есть несколько предобученных вариантов, которые существенно отличаются количеством параметров в сети, то есть напрямую влияет на производительность.

Ноутбук со всеми вычислениями и кодом тут

Мною была выбрана модель yolov8m.pt как "средний" вариант, в котором хорошо сбалансированы качество и производительность.

Дообучение производилось на 40-ка эпохах, при этом лучшая модель получилась уже в районе 15-20 эпох, потому что дальше началось переобучение

Пример результата работы данной сети для изображений:

Для видео:

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

Архитектура	mAP@50	Скорость (1 изображение), сек
Faster R-CNN	0.644	0.195
YOLOv8	0.339	0.063

Таким образом, если наиболее важна точность работы, то стоит выбирать двуступенчатые модели, такие как Faster R-CNN.

Если же важна скорость работы (например, работа с видео в режиме реального времени), то отличным вариантом будут одноступенчатые модели, например, из семейства YOLO