Небольшая предыстория. У нас в компании мы решили сделать коробочный проект по распознаванию номеров (для открытия шлагбаумов, учёта проезжающих машин и т.д. и т.п.). В качестве основы взяли одноплатный Jetson Nano.

В качестве эксперимента перенесли распознавание номеров, написанное на Python и работающее в облаке, на Jetson Nano. Результаты были неудовлетворительные - почти полная загрузка CPU по причине того, что использующийся для препроцессинга OpenCV работал на CPU и постпроцессинг модели Yolo также был на CPU.

К отрицательным моментам можно также отнести:

• Декодинг на CPU вместо аппаратного декодера, имеющегося у Jetson Nano

• Постоянно перекопирующуюся информацию из памяти GPU в память CPU (Хотя Jetson Nano делит обычную оперативную память между GPU и CPU).

• Ожидание GPU - 128 CUDA ядер просто сидели без дела пока 4 ARM ядра CPU были нагружены под завязку.

Решено было всё это переписывать на C++ и по возможности писать максимально оптимизированный код, близко приближенный к аппаратной платформе Jetson Nano, да и вообще завязанный на программных решениях от Nvidia.

На нашу удачу Nvidia в помощь разработчикам сделала библиотеку со всякими вспомогательными вещами для Jetson Nano. Благодаря ей можно инициализировать класс декодера видео просто указав URL RTSP потока, а получить в ответ очередь с кадрами прямо в CUDA видео-памяти, раскодированные аппаратным декодером.

// метод для создания декодера
gstDecoder* gstDecoder::Create( const videoOptions& options )
// в output записывает адрес кадра в видеопамяти
bool gstDecoder::Capture( void** output,
                         imageFormat format,
                         uint64_t timeout )

Препроцессинг

Кадр мы получили, но что делать дальше? (Yolov5s принимает нормализованное транспонированное из HWC в CHW изображение размером 640 на 640.) Необходим препроцессинг изображения:

1. Изменение размеров , добавление рамки ( для сохранения пропорций при размере 640 на 640, к тому же саму Yolo обучали с рамкой)

2. Перевод каждого пикселя из char во float32

3. Нормализация ( что бы значения каждого пикселя были от 0 до 1.0 включительно)

4. Транспонирование HWC => CHW

Для решения этой задачи мы решили использовать библиотеку Nvidia Perfomance Primitives или просто NPP. Кратко опишу как она работает. Например, нам необходимо изменить размер трехканального изображения (RGB), каждый пиксель которого хранится во float32, тогда берем функцию nppiResize_32f_C3R() . Название функции можно расшифровать так:

• 32f в названии говорит что каждый пиксель изображения хранится во float32

• C3 — что в изображении есть 3 канала

• R указывает, что функция работает с указанием ROI

Аргументы у функции следующие :

• const Npp8u *pSrc — указатель на начало исходного изображения

• int nSrcStep — шаг изображения по ширине в байтах

• NppiSize oSrcSize — ширина и высота исходного изображения

• NppiRect oSrcRectROI — прямоугольник ROI исходного изображения в формате {смещение по ширине, смещение по высоте, ширина, высота }

• Npp8u *pDst — указатель на начало нового изображения с измененными размерами

• int nDstStep — шаг нового изображения по ширине в байтах

• NppiSize oDstSize — ширина и высота нового изображения

• NppiRect oDstRectROI — прямоугольник ROI нового изображения

• int eInterpolation — интерполяция

С NPP достаточно непросто — если ошибся с аргументами то получил либо кашу вместо новой картинки, либо ненулевой статус ответа при выходе. С первым придётся сидеть и думать где же допустил ошибку. Второе проще — enum с ошибками на все случаи жизни точно подскажет что не так.

Для добавления рассчитанной рамки мы использовали nppiCopyConstBorder_32f_C3R, для нормализации изображения — nppiDivC_32f_C3IR, с транспонированием пару дней пришлось подумать как его лучше сделать , идеально подошла nppiCopy_32f_C3P3R:

// Делим на три слоя и последовательно копируем сразу в буфер Yolov5
// по смещению размером с один слой,
// тем самым транспонируя в нужный вид (CHW).
NppStatus st;
float * const inputArr[3] {this->yoloInput,
                 					 this->yoloInput + YOLO_PLANAR_OFFSET,
                 					 this->yoloInput + (YOLO_PLANAR_OFFSET * 2)};

int planarStep = YOLO_SIZE * 1 * sizeof(IMAGE_TYPE);

st = nppiCopy_32f_C3P3R( (Npp32f*)this->yoloBuffer,
                        	this->yoloConstStep,
                        	inputArr,
                        	planarStep,
                        	yoloSize );
return st;

Я не описал перевод изображения во float32, явно мы его не делали, переводом в float32 занимается библиотека сразу при захвате кадра :

// format указывает, в каком формате мы хотим получить изображение.
bool gstDecoder::Capture( void** output,
                         imageFormat format,
                         uint64_t timeout )

По поводу производительности такого препроцессинга на NPP — весь конвейер от захвата до записи в буффер TensorRT занимает на Jetson Nano 40 миллисекунд — это 25 кадров в секунду, которые приходят с камеры, замеры отдельно каждой стадии показывали примерно 10 - 15 миллисекунд.

Дальше происходит инференс Yolov5s на TensorRT , Yolov5 натренирована на меньшее количество классов чем стандартная, инференс занимает примерно 140-160 миллисекунд.

Постпроцессинг

На выходе Yolov5 мы получаем:

• 25200 координат прямоугольников предположительно распознанных объектов

• 25200 значений уверенности для этих прямоугольников

• (25200 * Количество классов) значений уверенности классов

  В коде результаты представлены вот так:

// x,y - координаты центра прямоугольника
// w,h - ширина, высота
// s - уверенность
// cs - уверенность в каждом классе 
// NUM_CLS_SCORES - количество классов Yolov5

struct YoloOut 
{
float x;
float y;
float w;
float h;
float s;
float cs[NUM_CLS_SCORES];
};
// Результат работы - 25200 YoloOut

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

Для отсеивания есть специальный алгоритм «Подавления без максимума» (англ. Non-Maximum suppression или просто NMS). В open source есть готовые имплементации этого алгоритма, например в OpenCV - cv::dnn::NMSBoxes(), но они работают на CPU и поэтому нам не подходили, ведь сначала бы понадобилось переместить результаты на CPU часть оперативки (TensorRT никак не хотел использовать память аллоцированную сразу и для CPU и для GPU), а потом ещё всё и отфильтровать с помощью CPU. Поэтому было решено сделать свой NMS на CUDA, что бы работало быстро и не занимало CPU.

Пришлось на ходу вникать как работает CUDA — сетка, блоки, нити, параллелизм , ограничения, квалификаторы и т.д. Для 25200 прямоугольников было придумано сделать ещё 25200 bool значений, в которых отражать отсеян элемент или нет.

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

// Разделим-скопируем отдельно прямоугольники, уверенность и класс объекта 
//по разным местам (наверное самое не оптимизированное место
// с точки зрения памяти)
__global__
void splitBoxes( const YoloOut * yoloOut,
                	Box * box,
                objConf_t * objectConfidence,
                ClsConf * class)

// Самая простая параллельность — 25200 блоков в сетке,
// по одной нити в блоке, в качестве аргументов
// указатели на начало памяти
  
splitBoxes <<<25200,1,1>>> (	this->yoloOut,
                            this->devBoxes,
                            this->objC,
                            this->clsConf);


// фильтрация по уверенности — если ниже порога, то ставим false
__global__
void filterObjConf( const objConf_t * objectConfidence,
                   boxStatus_t * boxStatus);

// точно также параллелим
filterObjConf <<<25200,1,1>>> (	this->objC,
                               this->boxStatus);

В чем удобство CUDA на мой взгляд — достаточно указать как тебе нужно распараллелить выполнение, а остальное CUDA берет на себя: распределяет вычисление 25200 элементов между 128 CUDA ядрами Jetson Nano.

Перед NMS помимо фильтрации по порогу уверенности необходимо умножить общую уверенность на максимальную уверенность класса и сохранить id класса с максимальной уверенностью:

// ci - сохраняем отдельно id класса, c максимумом уверенности
__global__
void mulObjAndCls( objConf_t * oc,
                  const ClsConf * cs,
                  clsIdx_t * ci,
                  const boxStatus_t * bs) // уже откинутые элементы не берем
  
  
mulObjAndCls <<<25200,1,1>>> (	this->objC,
                              this->clsConf,
                              this->clsId,
                              this->boxStatus);

После фильтрации по порогу уверенности остается не больше 1000 элементов (обычно около 100), которые нужно компактно вытащить из массива 25200 элементов YoloOut по 25200 bool ключам. Свой эффективный алгоритм параллельного сжатия массива на CUDA мы бы не потянули, поэтому для написания NMS нам очень помогла библиотека Thrust, с помощью неё сжали результаты при помощи thrust::copy_if(), работает действительно быстро — замеры показали 1-3 миллисекунды.

thrust::copy_if(objCTh,
                objCTh + YOLO_OUT_SIZE,
                boxStatusTh,
                this->objCThResult.begin(),
                is_true());

После сжатия идет сам алгоритм NMS, все оставшиеся прямоугольники попарно сравниваются, если превышен порог отношения пересечений и объединений (Intersection over Union, IoU), ставится флаг о непригодности, затем снова с помощью thrust::copy_if() происходит окончательное сжатие и запись результатов.

#define BLOCKSIZE 32
normIteratorDevPtr<Box> boxesR = thrust::copy_if(devBoxesTh,
                                                 devBoxesTh + YOLO_OUT_SIZE,
                                                 boxStatusTh,
                                                 this->boxesThResult.begin(),
                                                 is_true());

// resultSize это количество отфильтрованных прямоугольников до NMS
int resultSize = boxesR - boxesThResult.begin();

dim3 gridSize(int(ceil(float(resultSize)/BLOCKSIZE)),
              int(ceil(float(resultSize)/BLOCKSIZE)),1);

dim3 blockSize(BLOCKSIZE, BLOCKSIZE, 1);

_nms <<< gridSize, blockSize >>>(boxesRawPtr, boxStRawPtr);

normIteratorDevPtr<Box> resultNms = 
  thrust::copy_if(this->boxesThResult.begin(),
  this->boxesThResult.end(),
  this->boxStatusNmsResult.begin(),
  this->boxesAfterNmsResult.begin(),
  is_true());

resultSize = resultNms - this->boxesAfterNmsResult.begin();

Заключение

При такой оптимизации нагрузка на CPU снизилась до 10-15 %, а GPU используется на 99%.

Спасибо за внимание.