Работа с движками TensorRT в Windows средствами c++

Подходящий пример для рассмотрения этого вопроса лежит в sampleUffFasterRCNN (для версии TensorRT-7.0.0.11). Конечно там показано, как работать с моделью Uff, но в функции начальной загрузки build показано, как engine загружается из файла.  

std::vector<char> trtModelStream;

size_t size{0};

std::ifstream file(mParams.loadEngine, std::ios::binary);

if (file.good())

{

    file.seekg(0, file.end);

    size = file.tellg();

    file.seekg(0, file.beg);

    trtModelStream.resize(size);

    file.read(trtModelStream.data(), size);

    file.close();

}

После чего идет довольно простая инициализация:

IRuntime* infer = nvinfer1::createInferRuntime(gLogger);

if (mParams.dlaCore >= 0)

{

    infer->setDLACore(mParams.dlaCore);

}

mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(

    infer->deserializeCudaEngine(trtModelStream.data(), size, nullptr), samplesCommon::InferDeleter());

 

infer->destroy();

gLogInfo << "TRT Engine loaded from: " << mParams.loadEngine << endl;

if (!mEngine)

{

    return false;

}

else

{

    return true;

}

 Настройка модели

Конечно, в большинстве случаев мы заранее знаем, какие выходные и входные слои есть у движка, но если нет, то информацию об этом подскажет функция getInputOutputNames(). Инициализация описывается в функции initializeSampleParams(). Если работа идет только с движком, то важными будут следующие параметры:

1. Имена входных и выходных слоев
2. Имена объектов
3. Остальные параметры при тестировании можно не менять, но они читаются и понимаются без особых усилий

Рассмотрим настройки для двух моделей - для детектирования объектов FasterRCNN и классификации объектов с помощью resnet18.

Детектирование объектов FasterRCNN настраивается следующим образом. Типы входного и выходных слоев:

params.inputNodeName = "input_image";//faster rcnn

params.outputClsName = "dense_class_td/Softmax";

params.outputRegName = "dense_regress_td/BiasAdd";

params.outputProposalName = "proposal";

Классы объектов:

params.classNames.push_back("Automobile");

params.classNames.push_back("Bicycle");

params.classNames.push_back("Person");

params.classNames.push_back("Roadsign");

params.classNames.push_back("background");

Тут на самом деле нужно понимать, что названия особой роли не влияют. По сути внутри engine не хранятся имена, а соответствия индексов и имен нужно делать самому. 

Классификация образов с помощью resnet18 выглядит проще:

params.inputNodeName = "input_1"; //resnet18 classification

params.outputClsName = "predictions/Softmax";

и, если у вас всего 2 объекта, то:

params.classNames.push_back("val0");

params.classNames.push_back("val1");

Безо всякого дополнительного класса

 Распознавание

Распознавание реализуется в функции infer() и состоит из следующих этапов:

1. Создание буфера (если вы предполагаете распознавание не один раз, то создание буфера выносится в инициализацию)
2. Чтение данных (изображений)
3. Перенос данных на устройство (в видеопамять)
4. Распознавание данных
5. Получение результата

1. Создание буфера

Я не буду здесь вдаваться в подробности исходного кода, который находится за пределами рассматриваемого примера. Поэтому в данном случае для нас создание буфера выглядит очень простым:

samplesCommon::BufferManager buffers(mEngine, mParams.batchSize);

auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());

2. Чтение данных

В примере за это отвечает функция processInput и работает она с форматом ppm, но чего-то сложного в заполнении данных нет - это обычные 3 канала RGB. И вы легко можете заполнить их из того же изображения OpenCV. Например, ниже показано заполнение из OpenCV 1 канального Grayscale изображения (этого нет в описываемом примере):

float* hostDataBuffer = static_cast<float*>(buffers->getHostBuffer(layers_info[recog_mode].input));

float pixelMean[3]{ 103.939, 116.779, 123.68 };

 

for (int c = 0; c < 3; ++c) // 3 канала

{

       for(int i = 0; i < image.rows; i++)

             for(int j = 0; j < image.cols; j++)

                    hostDataBuffer[c * image.rows*image.cols + j + i*image.cols] = float(image.at<uchar>(i, j)) - pixelMean[c];

}

3. Перенос данных на устройство

buffers.copyInputToDevice()

 

4. Распознавание данных

status = context->execute(mParams.batchSize, buffers.getDeviceBindings().data());

5. Получение результата

Сначала данные копируются в обычную память с помощью buffers.copyOutputToHost(). А затем в verifyOutput() результат просматривается. И здесь рассмотрим опять 2 примера, как выше

1) детектирование объектов на изображении. Сначала происходит получение данных с выходных слоев:

const float* out_class = static_cast<const float*>(buffers.getHostBuffer(mParams.outputClsName));

const float* out_reg = static_cast<const float*>(buffers.getHostBuffer(mParams.outputRegName));

const float* out_proposal = static_cast<const float*>(buffers.getHostBuffer(mParams.outputProposalName));

Ну а заполнение данных о расположении объектов находится в функции  batch_inverse_transform_classifier(). 

if (x2 > x1 && y2 > y1)

{

    pred_boxes.push_back(x1);

    pred_boxes.push_back(y1);

    pred_boxes.push_back(x2);

    pred_boxes.push_back(y2);

    pred_probs.push_back(classifier_cls[n * roi_num_per_img * mParams.outputClassSize + max_idx

        + i * mParams.outputClassSize]);

    pred_cls_ids.push_back(max_idx);

    ++box_num;

}

2) для классификации объектов все гораздо проще. Из функции verifyOutput()  можно выйти сразу после вызова

const float* out_class = static_cast<const float*>(buffers.getHostBuffer(mParams.outputClsName));

Поскольку  out_class - это массив вероятностей всех классов. Наибольшая вероятность покажет наш искомый объект

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