背景
Hello，2020年底我安装配置好了TensorRT7，写了一篇文章总结了在Widnows10系统中如何配置与运行代码实现TensorRT开发环境的搭建。文章中详细介绍了配置的流程与步骤，文章的链接如下：
五分钟搞定VS2017 TensorRT环境搭建
当时还录了一个视频，我上传到了B站，觉得看文章麻烦，就看B站视频吧！地址在这里：
代码语言：javascript复制https://www.bilibili.com/video/BV1Bf4y167Ty今天给大家分享一下，我跑的第一个Pytorch TensorRT的模型训练与部署例子。不久以前我写过一个系列文章叫做《轻松学Pytorch系列》，其中有一篇CNN入门的文章，是讲如何通过CNN训练mnist数据集，然后导出模型ONNX格式，在OpenCV DNN中调用的。今天我就还继续用我导出的ONNX模型，实现它在TensorRT7中的调用，完成一个TensorRT版本的手写数字识别模型的部署。
配置
要想运行这个例子，还需要配置一波开发环境，在之前的配置的基础上，分别倒入安装好的CUDA的include目录，lib目录，然后把lib目录里面的*.lib文件统统扔到链接器中去。这样就完成了开发配置。我的CUDA安装路径如下：
代码语言：javascript复制C:Program FilesNVIDIA GPU Computing ToolkitCUDAv10.0
include目录为：
C:Program FilesNVIDIA GPU Computing ToolkitCUDAv10.0include
Lib目录为：
C:Program FilesNVIDIA GPU Computing ToolkitCUDAv10.0lib推理与演示
TensorRT的加载模型执行推理的步骤基本上跟OpenVINO与OpenCV DNN很相似，唯一区别的地方在于使用tensorRT做推理，首先需要把数据从内存搬到显存，处理完之后再重新搬回内存，然后解析输出。基本步骤与代码如下：创建网络
代码语言：javascript复制IBuilder* builder = createInferBuilder(gLogger);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
auto parser = nvonnxparser::createParser(*network, gLogger);解析ONNX格式模型文件
代码语言：javascript复制// 解析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");创建推理引擎
代码语言：javascript复制// 创建推理引擎
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 20);
config->setFlag(nvinfer1::BuilderFlag::kFP16);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
IExecutionContext *context = engine->createExecutionContext();获取输入输出格式与名称
代码语言：javascript复制// 获取输入与输出名称，格式
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);设置输入数据
代码语言：javascript复制// 预处理输入数据
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;创建输入/输出 显存缓冲区
代码语言：javascript复制// 创建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));执行推理
代码语言：javascript复制// 内存到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);解析输出与打印
代码语言：javascript复制// 解析输出
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运行结果如下：
