Около года назад компания Intel Movidius выпустила устройство для эффективного инференса сверточных нейросетей — Movidius Neural Compute Stick (NCS). Это устройство позволяет использовать нейросети для распознавания или детектирования объектов в условиях ограниченного энергопотребления, в том числе в задачах робототехники. NCS имеет USB-интерфейс и потребляет не более 1 ватта. В этой статье я расскажу об опыте использования NCS с Raspberry Pi для задачи обнаружения лиц в видео, включая как обучение Mobilenet-SSD детектора, так и его запуск на Raspberry.

Весь код можно найти в моих двух репозиториях: обучение детектора и демо с обнаружением лиц.



В своей первой статье я уже писал про обнаружение лиц с помощью NCS: тогда речь шла о YOLOv2 детекторе, который я конвертировал из формата Darknet в формат Caffe, а затем запускал на NCS. Процесс конвертирования оказался нетривиальным: поскольку эти два формата по-разному задают последний слой детектора, выход нейросети приходилось парсить отдельно, на CPU, с помощью куска кода из Darknet. Кроме того, этот детектор не удовлетворил меня как по скорости (до 5.1 FPS на моем ноутбуке), так и по точности — позже я убедился, что из-за чувствительности к качеству изображения от него сложно получить хороший результат на Raspberry Pi.

В итоге я решил просто обучить свой детектор. Выбор пал на SSD детектор с Mobilenet энкодером: легковесные свертки Mobilenet позволяют добиться высокой скорости без особых потерь в качестве, а сам SSD детектор не уступает YOLO и работает на NCS из коробки.

Обучение детектора

Архитектура

Код для обучения детектора расположен здесь.

Я решил воспользоваться готовым Mobilenet-SSD детектором, обученным на PASCAL VOC0712, и дообучить его на обнаружение лиц. Во-первых, это очень помогает обучать сетку быстрее, а во-вторых, не придется изобретать велосипед.

Исходный проект включал скрипт gen.py, который буквально собирал .prototxt файл модели, подставляя входные параметры. Я перенес его в свой проект, немного расширив функционал. Этот скрипт позволяет сгенерировать четыре типа конфигурационных файлов:

• train: на входе — обучающая LMDB база, на выходе — слой с подсчетом функции потерь и ее градиентов, есть BatchNorm
• test: на входе — тестовая LMDB база, на выходе — слой с подсчетом качества (mean average precision), есть BatchNorm
• deploy: на входе — изображение, на выходе — слой с предсказаниями, BatchNorm отсутствует
• deploy_bn: на входе — изображение, на выходе — слой с предсказаниями, есть BatchNorm

Последний вариант я добавил позже, чтобы в скриптах можно было загружать и преобразовывать сетку с BatchNorm, не трогая LMDB базы — иначе при отсутствии базы ничего не работало. (Вообще, мне кажется странным, что в Caffe источник данных задается в архитектуре сети — это как минимум не очень практично).


Архитектуру сети я слегка подкорректировал. Список изменений:

• Очевидно, число классов сменилось на 1 (не считая фона).
• Ограничения на соотношение сторон вырезаемых патчей при обучении: изменились с $$ на $$ (я решил немного упростить задачу и не обучаться на слишком растянутых картинках).
• Из дефолтных рамок остались только квадратные, по две на каждую ячейку. Их размеры я сильно уменьшил, поскольку лица существенно меньше, чем объекты в классической задаче детектирования объектов.

Caffe рассчитывает размеры дефолтных рамок так: имея минимальный размер рамки $$ и максимальный $$, она создает маленькую и большую рамки с размерами $$ и $$. Поскольку мне хотелось детектировать как можно более мелкие лица, я рассчитал полный stride для каждого слоя детекций и приравнял минимальный размер рамки к нему. При таких параметрах маленькие дефолтные рамки будут располагаться вплотную друг к другу и не будут пересекаться. Так у нас хотя бы есть гарантия, что пересечение с объектом будет существовать для какой-то рамки. Максимальный размер я установил вдвое больше. Для слоев conv16_2, conv17_2 я выставил размеры на глаз, одинаковыми. Таким образом, $$ для всех слоев составили: $$

Данные

Я использовал два датасета: WIDER Face и FDDB. WIDER содержит много картинок с очень мелкими и размытыми лицами, а FDDB больше тяготеет к крупным изображениям лиц (и на порядок меньше, чем WIDER). В них слегка различается формат аннотирования, но это уже детали.

Для обучения я использовал не все данные: я выкинул слишком маленькие лица (меньше шести пикселей или меньше 2% ширины изображения), выкинул все изображения с соотношением сторон меньше 0.5 или больше 2, выкинул все изображения, помеченные как «размытые» в датасете WIDER, поскольку они соответствовали по большей части совсем мелким лицам, и мне надо было хоть как-то выровнять соотношение мелких и крупных лиц. После этого я сделал все рамки квадратными, расширив наименьшую сторону: я решил, что меня не очень интересуют пропорции лица, а задача для нейросети немного упрощается. Также я выкинул все черно-белые картинки, которых было немного, и на которых скрипт сборки базы данных падает.

Чтобы использовать их для обучения и тестирования, надо собрать из них LMDB базу. Как это делается:

• Для каждого изображения создается разметка в .xml формате.
• Создается файл train.txt со строками вида "path/to/image.png path/to/labels.xml", такой же создается для test.
• Создается файл test_name_size.txt со строками вида "test_image_name height width"
• Создается файл labelmap.prototxt с числовыми соответствиями меткам

Запускается скрипт ssd-caffe/scripts/create_annoset.py (пример из Makefile):

python3 /opt/movidius/ssd-caffe/scripts/create_annoset.py --anno-type=detection --label-map-file=$(wider_dir)/labelmap.prototxt --min-dim=0 --max-dim=0 --resize-width=0 --resize-height=0 --check-label --encode-type=jpg --encoded --redo $(wider_dir) $(wider_dir)/trainval.txt $(wider_dir)/WIDER_train/lmdb/wider_train_lmdb ./data

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

После этого можно приступать к обучению.

Обучение

Код для обучения модели можно найти в моем Colab Notebook.

Обучение я производил в Google Colaboratory, поскольку мой ноутбук едва справлялся с тестированием сетки, а на обучении вообще зависал. Colaboratory позволила мне обучить сетку достаточно быстро и бесплатно. Подвох лишь в том, что мне пришлось написать скрипт компиляции SSD-Caffe для Colaboratory (включающий такие странные вещи, как перекомпиляцию boost и правку исходников), который выполняется порядка 40 минут. Подробнее можно узнать в моей предыдущей публикации.

Есть у Colaboratory еще одна особенность: после 12 часов машина умирает, безвозвратно стирая все данные. Лучший способ избежать потери данных — это монтировать в систему свой гугл диск и сохранять в него веса сети каждые 500-1000 итераций обучения.

Что касается моего детектора, за одну сессию в Colaboratory он успевал отучиться 4500 итераций, и полностью обучался за две сессии.

Качество предсказаний (mean average precision) на выделенном мной тестовом наборе данных (слитые WIDER и FDDB с ограничениями, перечисленными ранее) составило порядка 0.87 для лучшей модели. Для измерения mAP на сохраненных во время обучения весах есть скрипт scripts/plot_map.py.

Работа детектора на (очень странном) примере из датасета:

Запуск на NCS

Демо-программа с обнаружением лиц находится здесь.

Чтобы скомпилировать нейросеть для Neural Compute Stick, нужен Movidius NCSDK: он содержит утилиты для компиляции и профилирования нейросетей, а также C++ и Python API. Стоит заметить, что недавно была выпущена вторая версия, не совместимая с первой: все функции API были зачем-то переименованы, поменялся внутренний формат нейросеток, также добавились FIFO для взаимодействия с NCS и (наконец-то) появилось автоматическое преобразование из float 32 bit в float 16 bit, чего так не хватало в C++. Все свои проекты я обновил до второй версии, но оставил пару костылей для совместимости с первой.

После обучения детектора стоит слить BatchNorm слои с соседними свертками для ускорения нейросети. Этим занимается скрипт merge_bn.py отсюда, который я тоже позаимствовал из проекта Mobilenet-SSD.

Затем необходимо вызвать утилиту mvNCCompile, например:

mvNCCompile -s 12 -o graph_ssd -w ssd-face.caffemodel ssd-face.prototxt

В Makefile проекта для этого есть цель graph_ssd. Полученный файл graph_ssd является описанием нейросети в формате, понятном NCS.

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

• Получить дескриптор устройства по порядковому номеру
• Открыть устройство
• Считать скомпилированный файл нейросети в буфер (как бинарный файл)
• Создать пустой граф вычислений для NCS
• Разместить граф на устройстве, используя данные из файла, и выделить для него FIFO на input/output; буфер с содержанием файла теперь можно освободить
• Запуск детектора:
  ??
  • Получить изображение с камеры (или из любого другого источника)
  • Обработать его: отмасштабировать до нужного размера, преобразовать в float32 и привести к диапазону [-1,1]
  • Загрузить изображение на устройство и запросить инференс
  • Запросить результат (программа заблокируется до момента получения результата)
  • Распарсить результат, выделить рамки объектов (о формате — далее)
  • Вывести изображение с предсказаниями
  
  ??
• Освободить все ресурсы: удалить FIFO и граф вычислений, закрыть устройство и удалить его дескриптор

Практически для каждого действия с NCS есть своя отдельная функция, и в C++ выглядит это весьма громоздко, при этом приходится внимательно следить за освобождением всех ресурсов. Чтобы не нагружать код, я создал класс-обертку для работы с NCS. В нем вся работа по инициализации спрятана в конструктор и функцию load_file, а по освобождению ресурсов — в деструктор, и работа с NCS сводится к вызову 2-3 методов класса. К тому же, есть удобная функция для объяснения возникших ошибок.

Создаем обертку, передавая в конструктор размер входа и размер выхода (число элементов):

NCSWrapper NCS(NETWORK_INPUT_SIZE*NETWORK_INPUT_SIZE*3, NETWORK_OUTPUT_SIZE);

Загружаем скомпилированный файл с нейросеткой, попутно инициализируя все, что нам необходимо:

if (!NCS.load_file("./models/face/graph_ssd"))
{
    NCS.print_error_code();
    return 0;
}

Преобразуем изображение в float32 (image — это cv::Mat в формате CV_32FC3) и загружаем на устройство:

if(!NCS.load_tensor_nowait((float*)image.data))
{
    NCS.print_error_code();
    break;
}

Получаем результат (result — это свободный float указатель, буфер результата поддерживается оберткой); до окончания вычислений программа блокируется:

if(!NCS.get_result(result))
{
    NCS.print_error_code();
    break;
}

На самом деле, в обертке есть и метод, который позволяет загрузить данные и получить результат одновременно: load_tensor((float*)image.data, result). Я отказался от его использования не просто так: используя отдельные методы, можно слегка ускорить выполнение кода. После загрузки изображения CPU будет простаивать до тех пор, пока не придет результат выполнения с NCS (в данном случае это порядка 100 мс), и в это время можно заняться полезной работой: считать новый кадр и преобразовать его, а также вывести на экран предыдущие детекции. Именно так и реализована демо-программа, в моем случае это немного увеличивает FPS. Можно пойти дальше и запускать обработку изображений и детектор лиц асинхронно в двух разных потоках — это действительно работает и позволяет еще немного ускориться, однако в демо-программе не реализовано.

Детектор в качестве результата возвращает float массив размера 7*(keep_top_k+1). Здесь keep_top_k — параметр, заданный в .prototxt файле модели и показывающий, сколько детекций (в порядке уменьшения уверенности) нужно вернуть. Этот параметр, а также параметр, отвечающий за фильтрацию детекций по минимальному значению уверенности, и параметры non maximum suppression можно настроить в .prototxt файле модели в самом последнем слое. Стоит заметить, что если Caffe возвращает столько детекций, сколько на изображении было найдено, то NCS всегда возвращает keep_top_k детекций, чтобы размер массива был постоянным.

Сам массив результата устроен так: если рассматривать его как матрицу с keep_top_k+1 строками и 7 столбцами, то в первой строке, в первом элементе будет число детекций, а начиная со второй строки будут сами детекции в формате "garbage, class_index, class_probability, x_min, y_min, x_max, y_max". Координаты указаны в диапазоне [0,1], поэтому их необходимо будет домножить на высоту/ширину изображения. В остальных элементах массива будет мусор. При этом non maximum suppression выполняется автоматически, еще до получения результата (похоже, прямо на NCS).

Особенности запуска на Raspberry Pi

Сама демо-программа может быть запущена как на обычном компьютере или ноутбуке с Ubuntu, так и на Raspberry Pi с Raspbian Stretch. Я использую Raspberry Pi 2 model B, но демо должно работать и на других моделях. Makefile проекта содержит две цели для переключения режима: make switch_desk для компьютера/ноутбука и make switch_rpi для Raspberry Pi. Принципиальная разница в коде программы заключается лишь в том, что в первом случае для чтения данных с камеры используется OpenCV, а во втором случае — библиотека RaspiCam. Для запуска демо на Raspberry необходимо скомпилировать и установить ее.

Теперь очень важный момент: установка NCSDK. Если следовать стандартным инструкциям установки на Raspberry Pi, ничем хорошим это не кончится: установщик попытается подтащить и скомпилировать SSD-Caffe и Tensorflow. Вместо этого NCSDK нужно скомпилировать в API-only режиме. В этом режиме будут доступны только C++ и Python API (то есть, невозможно будет компилировать и профилировать графы нейросетей). Это значит, что граф нейросети нужно сначала скомпилировать на обычном компьютере, а затем скопировать на Raspberry. Для удобства я добавил в репозиторию два скомпилированных файла, для YOLO и для SSD.

Еще один интересный момент — это чисто физическое подключение NCS к Raspberry. Казалось бы, несложно подключить ее к USB-разъему, но нужно помнить, что ее корпус при этом заблокирует остальные три разъема (он довольно здоровый, так как выполняет функцию радиатора). Самый простой выход — подключить ее через USB-кабель.

Также стоит иметь в виду, что скорость выполнения будет различаться для разных версий USB (для конкретно этой нейросетки: 102 ms для USB 3.0, 92 ms для USB 2.0).

Теперь насчет питания NCS. Согласно документации, потребляет она до 1 ватта (при 5 вольтах на USB разъеме это будет до 200 ma; для сравнения: камера Raspberry потребляет до 250 ma). При питании от обычного зарядного устройства на 5 вольт, 2 ампера все прекрасно работает. Однако при попытке подключить две или больше NCS к Raspberry могут возникнуть проблемы. В этом случае рекомендуют использовать USB-разветвитель с возможностью внешнего питания.

На Raspberry демо работает медленнее, чем на компьютере/ноутбуке: 7.2 FPS против 10.4 FPS. Связано это с несколькими факторами: во-первых, от вычислений на CPU избавиться невозможно, а выполняются они намного медленнее; во-вторых, сказывается скорость передачи данных (для USB 2.0).

Также для сравнения я попытался запустить на Raspberry YOLOv2 детектор лиц из своей первой статьи, но заработал он очень плохо: при скорости в 3.6 FPS он пропускает множество лиц даже на простых кадрах. Судя по всему, он очень чувствителен к параметрам входного изображения, качество которого в случае камеры Raspberry далеко от идеала. SSD работает намного стабильнее, хотя пришлось немного подкрутить параметры видео в настройках RapiCam. он тоже иногда пропускает лица на кадре, но делает это довольно редко. Для увеличения стабильности в реальных приложениях можно добавить простой centroid tracker.

К слову: то же самое можно воспроизвести и на Python, есть туториал на PyImageSearch (используется Mobilenet-SSD для задачи object detection).

Другие идеи

Также я испытал пару идей по ускорению самой нейросети:

Первая идея: можно оставить только детекции слоев conv11 и conv13, а все лишние слои удалить. Получится детектор, который детектирует только мелкие лица и работает немного быстрее. В целом, не стоит того.

Вторая идея была интересная, но не заработала: я попытался выбросить из нейросетки свертки с весами, близкими к нулю, рассчитывая, что она станет быстрее. Однако таких сверток оказалось немного, и их удаление только слегка замедлило нейросеть (единственная догадка: это связано с тем, что число каналов перестало быть степенью двойки).

Заключение

Об обнаружении лиц на Raspberry я задумался довольно давно, как о подзадаче моего робототехнического проекта. Мне не понравились классические детекторы по соотношению скорости и качества, и я решил попробовать нейросетевые методы, заодно испытав Neural Compute Stick, в результате чего появились два проекта на GitHub и три статьи на Хабре (включая текущую). В целом, результат меня устраивает — скорее всего, именно этот детектор я буду использовать в своем роботе (возможно, о нем будет еще одна статья). Стоит заметить, что мое решение может оказаться не оптимальным — все же, это учебный проект, выполненный отчасти из любопытства к NCS. Все же, надеюсь, что эта статья окажется кому-нибудь полезной.