Поделиться через

Обучение модели обнаружения объектов с помощью TensorFlow

  • Статья

Итак, мы завершили настройку TensorFlow, и готовы применить архитектуру YOLO для обучения модели обнаружения объектов. YOLO — это нейронная сеть, которая за один проход прогнозирует для изображения положение ограничивающих прямоугольников и вероятности классификации. Модели YOLO способны обрабатывать более 60 кадров в секунду, а значит, эта архитектура отлично подходит для распознавания объектов в видео. Дополнительные сведения о работе YOLO вы найдете здесь.

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

Для начала скачайте этот файл с примером YOLO, который содержит вспомогательные скрипты для быстрого начала работы.

У вас есть три варианта использования YOLO:

  1. Применение заранее обученных весовых коэффициентов для модели YOLO. В этом сценарии модель заранее проходит обучение по крупному набору данных с 80 классами (категориями) для таких типичных объектов, как автобус, человек, бутерброд. Если вас интересует заранее обученная модель YOLO в формате ONNX, вы можете скачать ее здесь. После скачивания переходите сразу к завершающему этапу руководства, где описана интеграция этой модели в приложение.

  2. Реализация обучения переносом из пользовательского набора данных. Обучение переносом — это один из методов обучения модели, в котором заранее обученная модель берется за отправную точку и подготавливается к выполнению схожего, но существенно отличного задания. При работе с этим руководством мы применим заранее обученные весовые коэффициенты YOLO с 80 классами для создания новой модели с 20 классами, обучив ее по набору данных VOC. Если вы предпочитаете использовать собственный набор данных с другими классами, воспользуйтесь этими инструкциями.

  3. Обучение модели YOLO "с нуля". Мы не рекомендуем применять такой сценарий, так как в нем очень трудно обеспечить схождение. В исходном исследовании YOLO пришлось применить обучение по darknet в imagenet (с сотнями тысяч фотографий), а затем также обучить всю сеть.

Чтобы применить обучение переносом по заранее обученным весовым коэффициентам YOLO и набору данных VOC, выполните следующие действия:

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

  1. В окне miniconda перейдите к каталогу с примером YOLO и выполните указанную ниже команду, с помощью которой устанавливаются необходимые пакеты pip для YOLO.

pip install -r requirements.txt

  1. Выполнение скрипта установки для скачивания данных и заранее обученных весовых коэффициентов

python setup.py

  1. Преобразуйте набор данных. Реализация описана в файле tools/voc2012.py. Этот формат основан на API обнаружения объектов tensorflow. Многие поля здесь не являются обязательными, но они заполнены для совместимости с официальным API.
python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split train \
  --output_file ./data/voc2012_train.tfrecord

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split val \
  --output_file ./data/voc2012_val.tfrecord
  1. Обучение модели. Выполните следующие команды:
python convert.py
python detect.py --image ./data/meme.jpg # Sanity check

python train.py \
	--dataset ./data/voc2012_train.tfrecord \
	--val_dataset ./data/voc2012_val.tfrecord \
	--classes ./data/voc2012.names \
	--num_classes 20 \
	--mode fit --transfer darknet \
	--batch_size 16 \
	--epochs 10 \
	--weights ./checkpoints/yolov3.tf \
	--weights_num_classes 80

Итак, у вас теперь есть готовая к работе модель с 20 классами.

Следующие шаги

Завершив создание модели TensorFlow, нужно преобразовать ее в формат ONNX для использования с API-интерфейсами Windows Machine Learning.