Обнаружение лиц на видео с помощью Movidius Neural Compute Stick

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

Весь код можно найти на GitHub.

Подробнее об NCS

Neural Compute Stick — это устройство, предназначенное для ускорения нейронных сетей (преимущественно свёрточных) на этапе применения (inference). Идея заключается в том, что NCS можно присоединить к роботу или дрону и запускать нейросети там, где для этого не хватает вычислительных ресурсов. К примеру, NCS можно подключить к Raspberry Pi.

Фреймфорк для этого устройства, он же NCSDK, включает API для Python и C++, а также несколько полезных утилит, позволяющих скомпилировать нейронную сеть в формат, который понимает NCS, измерить время, которое занимают вычисления на каждом слое и проверить работоспособность сети. В качестве исходных данных могут выступать предобученные нейронные сети в формате Caffe или TensorFlow.

Выбор модели

Отлично, мы хотим решить задачу обнаружения лиц (face detection). Есть две довольно популярные архитектуры нейросетей для задач обнаружения: это Fast-RCNN/Faster-RCNN и YOLO. Мне не хотелось на данном этапе обучать свою модель, поэтому решил поискать готовую.

Трудность заключается в том, что NCSDK поддерживает далеко не все возможности, доступные в Caffe и TensorFlow, поэтому произвольная архитектура может просто не скомпилироваться. Например, далеко не все типы слоев поддерживаются, и при этом архитектура должна иметь ровно один входной слой (полный список ограничений и поддерживаемых слоев для Caffe можно увидеть здесь). Первая модель для обнаружения лиц, которую мне удалось найти (Faster-RCNN), не удовлетворяла обоим требованиям.

Затем я наткнулся на обученную модель архитектуры YOLO. Проблема была лишь в том, что нейросеть была в формате Darknet, хотя сама архитектура выглядела подходящей для NCS, поэтому появилась идея конвертировать нейросеть в формат Caffe.

Конвертация модели

Для конвертации модели я решил использовать вот этот проект, позволяющий переходить между форматами Darknet, Pytorch и Caffe.

Я запускаю конвертер в контейнере Docker — это уловка, которая появилась из-за того, что версия Caffe, установленная NCSDK, не понравилась конвертеру, а трогать конфигурацию системы мне не хотелось:

sudo docker run -v `pwd`:/workspace/data \
	-u `id -u` -ti dlconverter:latest bash -c \
	"python ./pytorch-caffe-darknet-convert/darknet2caffe.py \
	./data/yolo-face.cfg ./data/yolo-face_final.weights \
	./data/yolo-face.prototxt ./data/yolo-face.caffemodel"

Получается кое-что интересное: модель конвертируется, но выдается предупреждение о том, что слои типа Crop, Dropout и Detection распознать не удалось, из-за чего конвертер их пропустил. Можно ли обойтись без этих слоев? Оказывается, можно. Если внимательно посмотреть на код Darknet, можно заметить следующее:

Слой типа Crop нужен только на этапе обучения. Он занимается тем, что расширяет выборку, поворачивая изображение на случайные углы и вырезая из него случайные фрагменты. На этапе применения он не потребуется.

С Dropout немного интереснее. Dropout слой тоже нужен в основном на этапе обучения (Dropout слой обнуляет выходы нейронов с вероятностью ) для того, чтобы избежать переобучения и повысить способности модели к обобщению. На этапе применения от него можно избавиться, но при этом необходимо масштабировать выходы нейронов так, чтобы матожидание значений на входах следующего слоя не изменилось, чтобы поведение модели сохранилось (разделить на ). Если вглядеться в код Darknet, то можно заметить, что Dropout слой не только обнуляет некоторые выходы, но и масштабирует все остальные, поэтому Dropout слой можно безболезненно удалить.

Что касается Detection слоя, то он находится последним и занимается тем, что переводит в более читаемый вид выходы с предпоследнего слоя, а также считает функцию потерь на этапе обучения. Функцию потерь нам считать не нужно, а вот перевод результата в читаемый вид пригодится. В итоге я решил просто использовать функцию для последнего слоя прямиком из Darknet (немного подредактировав ее), оттуда же взял функцию для NMS (Non maximum suppression — удаление избыточных ограничивающих рамок). Они находятся в файлах detection_layer.c и detection_layer.h.

Тут стоит сделать замечание о том, что делает предпоследний слой. В архитектуре YOLO (You only look once) изображение разбивается на блоки сеткой размера (в данном случае ), и для каждого блока предсказывается значений, где — число ограничивающих рамок для каждого блока, а — число классов. Сами значения представляют собой: координаты, ширину, высоту и значение уверенности для каждой из рамок (то есть, всего пять значений), а также вероятность нахождения объекта в этом блоке для каждого класса. Итого получается значений. Слой Detection только разделяет разные типы данных и приводит их в структурированный вид.

Осталась одна проблема: конвертер генерирует .prototxt файл с форматом входного слоя, который NCSDK не может разобрать. Различия исключительно декоративные: конвертер записывает размер входного слоя в формате:

input_dim: x
input_dim: y 
input_dim: z 
input_dim: w

А утилита mvNCCompile, которая должна компилировать нейронную сеть в понятный NCS файл, хочет видеть формат:

input_shape {
dim: x
dim: y
dim: z
dim: w
}

Python скрипт utils/fix_proto_input_format.py призван решить эту проблему (не самому же это делать).

Компиляция модели

Теперь, когда модель переведена в формат Caffe, можно ее скомпилировать. Делается это довольно просто:

mvNCCompile -s 12 -o graph -w yolo-face.caffemodel yolo-face-fix.prototxt

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

Предобработка изображений

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

Судя по коду демо для Darknet, перед загрузкой изображения его нужно сжать до размера (причем не заботясь о пропорциях), нормировать на отрезок каждый пиксель и инвертировать порядок каналов с BGR на RGB. Вообще, в Caffe и OpenCV стандартным считается вариант BGR, а в Darknet — RGB, однако конвертер ничего об этом не знает, и в итоге каналы все равно нужно инвертировать.

Обращение к NCS и загрузка данных

Тут стоит заметить, что я использую C++, а не Python, поскольку ориентируюсь на применение устройства в робототехнике и верю, что на C++ можно добиться большего быстродействия. Именно из-за этого появляются дополнительные сложности: граф вычислений получает на вход и выдает на выходе данные в формате fp16 (16-битные числа с плавающей точкой), реализации которых нет в C++ по умолчанию. В примерах NCSDK эта проблема решается использованием функций floattofp16 и fp16tofloat, выдранных из Numpy, поэтому я использую такое же решение.

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

• Вызвать mvncGetDeviceName, чтобы получить имя NCS
• Открыть устройство по имени с помощью mvncOpenDevice
• Загрузить содержимое файла graph в буфер (специальной функции для этого нет, нужно использовать свою)
• Разместить граф вычислений с помощью mvncAllocateGraph

Для загрузки данных и получения результата используются функции mvncLoadTensor и mvncGetResult соответственно. При этом нужно помнить про конвертацию данных и результата в fp16 и обратно.

Для прекращения работы с NCS нужно освободить ресурсы, отведенные под граф вычислений (mvncDeallocateGraph) и закрыть устройство (mvncCloseDevice).

Поскольку для взаимодействия с NCS нужно довольно много действий, я написал класс-обертку, у которого (помимо конструктора и деструктора) есть всего две функции: load_file для инициализации устройства и графа и load_tensor для загрузки данных и получения результата.

Профилировщик

В NCSDK есть полезная утилита, которая позволяет не только оценить быстродействие каждого слоя, но еще и создать схему графа вычислений, на которой отражены характеристики каждого элемента (кстати, передавать сами веса слоев при этом необязательно):

mvNCProfile yolo-face-fix.prototxt -w yolo-face.caffemodel -s 12

Что получилось в итоге

Итоговый результат выглядит следующим образом:

Исходная модель (.cfg и .weights) с помощью конвертера преобразуется в формат Caffe (.prototxt и .caffemodel), в файле .prototxt формат входов исправляется с помощью Python скрипта, после чего модель компилируется в файл graph — это цели convert и graph в Makefile.

В самой программе для каждого полученного кадра производится предобработка, перевод в формат fp16 и загрузка в граф вычислений. Полученный результат переводится из fp16 в формат float и передается в функцию, которая имитирует работу последнего Detection слоя. Затем применяется Non maximum suppression.

Демо гордо выдает 4.5 кадра в секунду — это маловато. Проблема, видимо, в том, что эта архитектура относится к тому типу, который затачивается больше на точность, а не на быстродействие. Скорость работы можно значительно повысить, если использовать «мобильные» архитектуры вроде Tiny YOLO — для этого придется искать новую модель или обучать свою. Тем не менее, этот пример показывает, что нейросеть в формате Darknet можно скомпилировать и запустить на Neural Compute Stick.