推理演示 | 八步助你搞定tensorRT C++ SDK调用！

OpenCV学堂

发布于 2021-03-12 16:28:12

6.9K2

发布于 2021-03-12 16:28:12

文章被收录于专栏：贾志刚-OpenCV学堂贾志刚-OpenCV学堂

背景

Hello，2020年底我安装配置好了TensorRT7，写了一篇文章总结了在Widnows10系统中如何配置与运行代码实现TensorRT开发环境的搭建。文章中详细介绍了配置的流程与步骤，文章的链接如下：

五分钟搞定VS2017+TensorRT环境搭建

当时还录了一个视频，我上传到了B站，觉得看文章麻烦，就看B站视频吧！地址在这里：

https://www.bilibili.com/video/BV1Bf4y167Ty

今天给大家分享一下，我跑的第一个Pytorch + TensorRT的模型训练与部署例子。不久以前我写过一个系列文章叫做《轻松学Pytorch系列》，其中有一篇CNN入门的文章，是讲如何通过CNN训练mnist数据集，然后导出模型ONNX格式，在OpenCV DNN中调用的。今天我就还继续用我导出的ONNX模型，实现它在TensorRT7中的调用，完成一个TensorRT版本的手写数字识别模型的部署。

配置

要想运行这个例子，还需要配置一波开发环境，在之前的配置的基础上，分别倒入安装好的CUDA的include目录，lib目录，然后把lib目录里面的*.lib文件统统扔到链接器中去。这样就完成了开发配置。我的CUDA安装路径如下：

C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.0
include目录为：
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.0\include
Lib目录为：
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.0\lib

推理与演示

TensorRT的加载模型执行推理的步骤基本上跟OpenVINO与OpenCV DNN很相似，唯一区别的地方在于使用tensorRT做推理，首先需要把数据从内存搬到显存，处理完之后再重新搬回内存，然后解析输出。基本步骤与代码如下：创建网络

IBuilder* builder = createInferBuilder(gLogger);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
auto parser = nvonnxparser::createParser(*network, gLogger);

解析ONNX格式模型文件

// 解析ONNX模型
parser->parseFromFile(onnx_filename.c_str(), 2);
for (int i = 0; i < parser->getNbErrors(); ++i)
{
    std::cout << parser->getError(i)->desc() << std::endl;
}
printf("tensorRT load onnx mnist model...\n");

创建推理引擎

// 创建推理引擎
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 20);
config->setFlag(nvinfer1::BuilderFlag::kFP16);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
IExecutionContext *context = engine->createExecutionContext();

获取输入输出格式与名称

// 获取输入与输出名称，格式
const char* input_blob_name = network->getInput(0)->getName();
const char* output_blob_name = network->getOutput(0)->getName();
printf("input_blob_name : %s \n", input_blob_name);
printf("output_blob_name : %s \n", output_blob_name);

const int inputH = network->getInput(0)->getDimensions().d[2];
const int inputW = network->getInput(0)->getDimensions().d[3];
printf("inputH : %d, inputW: %d \n", inputH, inputW);

设置输入数据

// 预处理输入数据
Mat image = imread("D:/images/9_99.png", IMREAD_GRAYSCALE);
imshow("输入图像", image);
Mat img2;
image.convertTo(img2, CV_32F);
img2 = (img2 / 255 - 0.5) / 0.5;

创建输入/输出显存缓冲区

// 创建GPU显存输入/输出缓冲区
void* buffers[2] = { NULL, NULL };
int nBatchSize = 1;
int nOutputSize = 10;
cudaMalloc(&buffers[0], nBatchSize * inputH * inputW * sizeof(float));
cudaMalloc(&buffers[1], nBatchSize * nOutputSize * sizeof(float));

// 创建cuda流
cudaStream_t stream;
cudaStreamCreate(&stream);
void *data = malloc(nBatchSize * inputH * inputW * sizeof(float));
memcpy(data, img2.ptr<float>(0), inputH * inputW * sizeof(float));

执行推理

// 内存到GPU显存
cudaMemcpyAsync(buffers[0], data, \
nBatchSize * inputH * inputW * sizeof(float), cudaMemcpyHostToDevice, stream);
std::cout << "start to infer image..." << std::endl;

// 推理
context->enqueueV2(buffers, stream, nullptr);

// 显存到内存
float prob[10];
cudaMemcpyAsync(prob, buffers[1], 1 * nOutputSize * sizeof(float), cudaMemcpyDeviceToHost, stream);

// 同步结束，释放资源
cudaStreamSynchronize(stream);
cudaStreamDestroy(stream);

解析输出与打印

// 解析输出
std::cout << "image inference finished!" << std::endl;
Mat result = Mat(1, 10, CV_32F, (float*)prob);
float max = result.at<float>(0, 0);
int index = 0;
for (int i = 0; i < 10; i++)
{
    if (max < result.at<float>(0,i)) {
        max = result.at<float>(0, i);
        index = i;
    }
}
std::cout << "predict digit: " << index << std::end

运行结果如下：