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

TensorRT позволяет в несколько раз ускорить работу моделей распознавания за счет настраивания на оптимальную форму сети для конкретной видеокарты, перевод в Float16 и Int8. В Windows средства TensorRT работаю только в C++, а иногда все-таки требуется, чтобы машинное обучение работало и в это операционной системе. Поэтому в данной статье показано, как работать с движком (engine) TensorRT в Windows. В подробности получения движка engine вдаваться не будем, предположим, что он у нас уже есть. В остальном, как вы увидите ниже, несмотря на кажущуюся сложность кода, работа напрямую с движком очень проста с точки зрения принципа.

 Загрузка engine

Подходящий пример для рассмотрения этого вопроса лежит в 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 не хранятся имена, а соответствия индексов и имен нужно делать самому. 

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

Классификация образов с помощью 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 — это массив вероятностей всех классов. Наибольшая вероятность покажет наш искомый объект

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

Please follow and like us: