Не так давно в свет вышло устройство 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.