Работа с движками 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(). Если работа идет только с движком, то важными будут следующие параметры:
- Имена входных и выходных слоев
- Имена объектов
- Остальные параметры при тестировании можно не менять, но они читаются и понимаются без особых усилий
Рассмотрим настройки для двух моделей - для детектирования объектов 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. Создание буфера
Я не буду здесь вдаваться в подробности исходного кода, который находится за пределами рассматриваемого примера. Поэтому в данном случае для нас создание буфера выглядит очень простым:
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 - это массив вероятностей всех классов. Наибольшая вероятность покажет наш искомый объект
В итоге можно сделать вывод о том, что не смотря на обилие кода, сама структура распознавания понятна и ее в буквальном смысле можно потрогать руками, настроив в своем продукте параметры таким образом, как вам удобно.