Pytorch导出onnx模型，C++转化为TensorRT并实现推理过程

前言

1. 本文为旨在实现整个Python导出PyTorch模型，C++转化为TensorRT并实现推理过程过程，只与模型推理，模型部署相关，不涉及模型训练。
2. 为突出整个部署过程而非具体模型本身，本文模型就采用最简单的分类模型，并且直接使用 torchvision.model 中的权重。检测、分割等其他模型在前后处理部分会有不同，但是模型本身的导出、转换和推理的过程基本是一致的。
3. 本文会先用 Pytorch 测试一个分类模型在一张测试图片上的结果，将其转换为 onnx 模型，再用 onnxruntime 测试结果，再用 C++ 将其转换为 TensorRT 模型，再测试推理结果。预期三者测试结果一致，则转换成功。
4. 如果想要测试速度，Python 可以使用 time.perf_counter()， C++ 可以使用 std::chrono::high_resolution_clock 。建议多测数据集中的一些图片，计算推理时间的均值和方差，而不是只测一张图片。

1 Pytorch模型推理测试导出onnx

这部分我们使用 torchvision 实例化一个简单的 ResNet50 分类模型，并将其导出为 onnx 模型。在这个过程中，我们还需要使用一张图片进行推理，并记录下 Python 模型的输出，方便我们后面到处 TensoRT 模型并进行推理时进行准确性的验证。

由于 torchvision 中的 resnet50 分类模型中是没有进行最后的 softmax 操作的，这里我们为了之后使用方便，自己新建一个类 ResNet50_wSoftmax 将后处理 softmax 添加到模型中一起导出。

这也是 pytorch 导出 onnx 模型的一个推荐的方式，就是将一些必要后处理添加到模型中一起导出，这样做有两个优点：

• 可以直接得到端到端的 onnx/tensorrt 模型，不必在外面再做后处理操作
• 再之后我们会将 onnx 模型转换为 tensorrt 模型，在转换过程中 tensorrt 会对我们的模型进行一些针对特定的 Nvidia GPU 的推理优化，我们将后处理一起合并到 onnx 模型中，可能可以使得一些算子操作再转换为 tensorrt 的过程中同样得到优化。

最终代码如下：

# export_onnx.py
import torch
import torchvision.models as models
import cv2
import numpy as npclass ResNet50_wSoftmax(torch.nn.Module):# 将softmax后处理合并到模型中，一起导出为onnxdef __init__(self):super().__init__()self.base_model = models.resnet50(pretrained=True)self.softmax = torch.nn.Softmax(dim=1)def forward(self, x):y = self.base_model(x)prob = self.softmax(y)return probdef preprocessing(img):# 预处理：BGR->RGB、归一化/除均值减标准差IMAGENET_MEAN = [0.485, 0.456, 0.406]IMAGENET_STD = [0.229, 0.224, 0.225]img = img[:, :, ::-1]img = cv2.resize(img, (224, 224))img = img / 255.0img = (img - IMAGENET_MEAN) / IMAGENET_STDimg = img.transpose(2, 0, 1).astype(np.float32)tensor_img = torch.from_numpy(img)[None]return tensor_imgif __name__ == '__main__':# model = models.resnet50(pretrained=True)image_path = 'test.jpg'img = cv2.imread(image_path)tensor_img = preprocessing(img)model = ResNet50_wSoftmax()   # 将后处理添加到模型中model.eval()pred = model(tensor_img)[0]max_idx = torch.argmax(pred)print(f"test_image: {image_path}, max_idx: {max_idx}, max_logit: {pred[max_idx].item()}")dummpy_input = torch.zeros(1, 3, 224, 224)  # onnx的导出需要指定一个输入，这里直接用上面的tenosr_img也可torch.onnx.export(model, dummpy_input, 'resnet50_wSoftmax.onnx',input_names=['image'],output_names=['predict'],opset_version=11,dynamic_axes={'image': {0: 'batch'}, 'predict': {0: 'batch'}}		# 注意这里指定batchsize是动态可变的)

执行结果会输出：

test_image: test.jpg, max_idx: 971, probability: 0.994541585445404

这些结果我们一会测试 onnx/tensorrt 模型时用于比对转换是否有误差。并得到一个 onnx 模型文件：classifier.onnx 。

2 onnxruntime推理测试

我们将刚刚得到的 classifier.onnx ，用 onnxruntime 来进行推理测试，看结果是否相同。

这里，我们就复用刚才测试 pytorch 模型时的预处理函数，整个 onnxruntime 推理测试代码如下：

import onnxruntime as ort
import numpy as np
import cv2
from export_onnx import preprocessingimage_path = 'test.jpg'
ort_session = ort.InferenceSession("classifier.onnx") # 创建一个推理sessionimg = cv2.imread(image_path)
input_img = preprocessing(img)[None]pred = ort_session.run(None, { 'image' : input_img } )[0][0]
max_idx = np.argmax(pred)
print(f"test_image: {image_path}, max_idx: {max_idx}, probability: {pred[max_idx]}")

输出：

test_image: test.jpg, max_idx: 971, probability: 0.994541585445404

可以看到，跟我们 pytorch 模型的测试结果是一致的。

3 C++ onnx模型转换为tensorrt模型

本部分重度参考自课程：tensorRT从零起步迈向高性能工业级部署（就业导向）

我们进行模型部署推理肯定是追求极致的推理速度，这时再用 Python 来进行转换和推理就不合适了，接下来我们就转战到 C++ 上，将onnx模型转换为tensorrt模型。

对于大部分深度学习部署的 C/C++ 的初学者而言，环境配置都是个老大难的问题。本身 C/C++ 的包管理就不如 Python 的 pip、conda 等来的直接方便，再加上各种 nvidia driver/cuda/cudnn/cuda-runtime 的各种版本不对齐的问题，包括笔者在内的许多萌新们初期总是会在环境配置遇到许多问题。但是本文关注的重点是整个模型转换和部署的过程，不可能花大篇幅再去介绍环境配置，将来有机会再单独写一篇介绍 Python/C++ 深度学习模型部署时环境配置的问题，这里就直接给出笔者使用的关键软硬件的版本号/型号。

GPU: RTX 3060ti 12GB

OS: ubuntu 18.04

gcc: 7.5

TensorRT: 8.x

CUDA: 11.2

cuDNN: 8.x

头文件

包含的头文件：

// tensorrt相关
#include <NvInfer.h>
#include <NvInferRuntime.h>// onnx解析器相关
#include <onnx-tensorrt/NvOnnxParser.h>// cuda_runtime相关
#include <cuda_runtime.h>// 常用头文件
#include <stdio.h>
#include <math.h>
#include <string>
#include <iostream>
#include <fstream>
#include <vector>
#include <memory>
#include <functional>
#include <unistd.h>
#include <chrono>// opencv
#include <opencv2/opencv.hpp>

logger类

首先我们要准备一个 logger 类，来打印构建 tensorrt 模型过程中的一些错误或警告。按照指定的严重性程度 (severity)，来打印信息。

inline const char* severity_string(nvinfer1::ILogger::Severity t) {switch (t) {case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";case nvinfer1::ILogger::Severity::kERROR: return "error";case nvinfer1::ILogger::Severity::kWARNING: return "warning";case nvinfer1::ILogger::Severity::kINFO: return "info";case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";default: return "unknown";}
}class TRTLogger : public nvinfer1::ILogger {
public:virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {if (severity <= Severity::kWARNING) {if (severity == Severity::kWARNING) printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);else if (severity == Severity::kERROR) printf("\031[33m%s: %s\033[0m\n", severity_string(severity), msg);else printf("%s: %s\n", severity_string(severity), msg);}}
};

build_model函数

build_model 函数，各步骤已在代码中添加注释：

bool build_model() {if (isFileExist( "classifier.trtmodel" )) {printf("classifier.trtmodel already exists.\n");return true;}TRTLogger logger;// 下面的builder, config, network是基本需要的组件// 形象的理解是你需要一个builder去build这个网络，网络自身有结构，这个结构可以有不同的配置nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);// 创建一个构建配置，指定TensorRT应该如何优化模型，tensorRT生成的模型只能在特定配置下运行nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();// 创建网络定义，其中createNetworkV2(1)表示采用显性batch size，新版tensorRT(>=7.0)时，不建议采用0非显性batch sizenvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);// onnx parser解析器来解析onnx模型auto parser = nvonnxparser::createParser(*network, logger);if (!parser->parseFromFile("classifier.onnx", 1)) {printf("Failed to parse classifier.onnx.\n");return false;}// 设置工作区大小printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);config->setMaxWorkspaceSize(1 << 28);// 需要通过profile来使得batchsize时动态可变的，这与我们之前导出onnx指定的动态batchsize是对应的int maxBatchSize = 10;auto profile = builder->createOptimizationProfile();auto input_tensor = network->getInput(0);auto input_dims = input_tensor->getDimensions();// 设置batchsize的最大/最小/最优值input_dims.d[0] = 1;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, input_dims);profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, input_dims);input_dims.d[0] = maxBatchSize;profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, input_dims);config->addOptimizationProfile(profile);// 开始构建tensorrt模型enginenvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);if (engine == nullptr) {printf("Build engine failed.\n");return false;}// 将构建好的tensorrt模型engine反序列化（保存成文件）nvinfer1::IHostMemory* model_data = engine->serialize();FILE* f = fopen("classifier.trtmodel", "wb");fwrite(model_data->data(), 1, model_data->size(), f);fclose(f);// 逆序destory掉指针model_data->destroy();engine->destroy();network->destroy();config->destroy();builder->destroy();printf("Build Done.\n");return true;
}

调用 build_model 函数成功后，我们会得到一个 classifier.trtmodel 文件。

make_nvshared

上面的实现有个比较不优雅的地方，对于我们创建的 builder、config 等指针，我们都需要一一进行 destroy，从而避免内存泄漏。实际上，这里我们可以通过共享指针，来实现自动释放。

shared_ptr<_T> make_nvshared(_T *ptr) {return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}

在这里指定一下释放内存的方式，之后就可以通过类似：

auto network = make_nvshared(builder->createNetworkV2(1));

这样的方式创建智能指针，他会自己 destroy 释放，这样最后几行 destory 就不用写了。

4 tensorrt模型推理测试

我们上一步已经成功将 onnx 模型导出为了 tensorrt 模型，现在我们用 tensorrt 模型来进行推理，看一下结果是否与之前 pytorch 和 onnx 推理的结果一致，如果一致，则模型转换成功。

load_file

load_file 函数用于加载我们的 tensorrt 模型：

vector<unsigned char> load_file(const string& file) {ifstream in(file, ios::in | ios::binary);if (!in.is_open()) return {};in.seekg(0, ios::end);size_t length = in.tellg();vector<uint8_t> data;if (length > 0) {in.seekg(0, ios::beg);data.resize(length);in.read((char*)&data[0], length);}in.close();return data;
}

inference

void inference(const string& image_path) {TRTLogger logger;// 加载模型auto engine_data = load_file("classifier.trtmodel");// 执行推理前，需要创建一个推理的runtime接口实例。与builer一样，runtime需要loggerauto runtime = make_nvshared(nvinfer1::createInferRuntime(logger));auto engine = make_nvshared(runtime->deserializeCudaEngine(engine_data.data(), engine_data.size()));if (engine == nullptr) {printf("Deserialize cuda engine failed.\n");runtime->destroy();return;}if (engine->getNbBindings() != 2) {printf("Must be single input, single Output, got %d output.\n", engine->getNbBindings() - 1);return;}// 创建CUDA流，以确定这个batch的推理是独立的cudaStream_t stream = nullptr;checkRuntime(cudaStreamCreate(&stream));auto execution_context = make_nvshared(engine->createExecutionContext());int input_batch = 1;int input_channel = 3;int input_height = 224;int input_width = 224;// 准备好input_data_host和input_data_device，分别表示内存中的数据指针和显存中的数据指针// 一会儿将预处理过的图像数据搬运到GPUint input_numel = input_batch * input_channel * input_height * input_width;float* input_data_host = nullptr;float* input_data_device = nullptr;checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));// 图片读取与预处理，与之前python中的预处理方式一致：// BGR->RGB、归一化/除均值减标准差float mean[] = {0.406, 0.456, 0.485};float std[] = {0.225, 0.224, 0.229};auto image = cv::imread(image_path);cv::resize(image, image, cv::Size(input_width, input_height));int image_area = image.cols * image.rows;unsigned char* pimage = image.data;float* phost_b = input_data_host + image_area * 0;float* phost_g = input_data_host + image_area * 1;float* phost_r = input_data_host + image_area * 2;for (int i=0; i<image_area; ++i, pimage += 3) {*phost_r++ = (pimage[0] / 255.0f - mean[0]) / std[0];*phost_g++ = (pimage[1] / 255.0f - mean[1]) / std[1];*phost_b++ = (pimage[2] / 255.0f - mean[2]) / std[2];}// 进行推理checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel *sizeof(float), cudaMemcpyHostToDevice, stream));const int num_classes = 1000;float output_data_host[num_classes];float* output_data_device = nullptr;checkRuntime(cudaMalloc(&output_data_device, sizeof(output_data_host)));auto input_dims = engine->getBindingDimensions(0);input_dims.d[0] = input_batch;execution_context->setBindingDimensions(0, input_dims);// 用一个指针数组bindings指定input和output在gpu中的指针。float* bindings[] = {input_data_device, output_data_device};bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);checkRuntime(cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream));checkRuntime(cudaStreamSynchronize(stream));float* prob = output_data_host;int predict_label = max_element(prob, prob + num_classes) - prob;float conf = prob[predict_label];printf("test_image: %s, max_idx: %d, probability: %f", image_path.c_str(), predict_label, conf);// 释放显存checkRuntime(cudaStreamDestroy(stream));checkRuntime(cudaFreeHost(input_data_host));checkRuntime(cudaFree(input_data_device));checkRuntime(cudaFree(output_data_device));
}

最终得到输出：

test_image: test.jpg, max_idx: 971, probability: 0.994527

与之前 pytorch 和 onnx 推理的结果基本一致，模型转换成功。

附录

给出完整的参考代码：https://github.com/Adenialzz/Hello-AIDeployment/tree/master/HAID/tensorrt/resnet