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

本文为旨在实现整个 Python导出PyTorch模型，C++转化为TensorRT并实现推理过程 过程，只与模型推理，模型部署相关，不涉及模型训练。
为突出整个部署过程而非具体模型本身，本文模型就采用最简单的分类模型，并且直接使用 torchvision.model 中的权重。检测、分割等其他模型在前后处理部分会有不同，但是模型本身的导出、转换和推理的过程基本是一致的。
本文会先用 Pytorch 测试一个分类模型在一张测试图片上的结果，将其转换为 onnx 模型，再用 onnxruntime 测试结果，再用 C++ 将其转换为 TensorRT 模型，再测试推理结果。预期三者测试结果一致，则转换成功。
如果想要测试速度，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 np
class ResNet50_wSoftmax(torch.nn.Module):
    # 将softmax后处理合并到模型中，一起导出为onnx
    def __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 prob
def 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.0
    img = (img - IMAGENET_MEAN) / IMAGENET_STD
    img = img.transpose(2, 0, 1).astype(np.float32)
    tensor_img = torch.from_numpy(img)[None]
    return tensor_img
if __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 preprocessing
image_path = 'test.jpg'
ort_session = ort.InferenceSession("classifier.onnx") # 创建一个推理session
img = 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 size
	nvinfer1::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模型engine
	nvinfer1::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需要logger
	auto 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，分别表示内存中的数据指针和显存中的数据指针
  // 一会儿将预处理过的图像数据搬运到GPU
	int 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
				首先，无论是python还是C++，我们都要写一个自己的plugin，一般的不支持op网上都有参考的版本，写plugin的参考教程：如何写好一个op
C++实现和使用过程：
下载TensorRT官网源码， 进入plugin可以看到许多已经定义好的插件可以参考，照着写就行，以ScatterND为例。
先写好自己的plugin，如名为：myplugin
注意：此处plugin代码里面不能包含REGISTER_TENSORRT_PLUGIN(MyPluginCreator)代码，不然会导致重复注册的问题。
				一文掌握Pytorch-onnx-tensorrt模型转换pytorch转onnx
pytorch转onnx
对于实际部署的需求，很多时候pytorch是不满足的，所以需要转成其他模型格式来加快推理。常用的就是onnx，onnx天然支持很多框架模型的转换，
如Pytorch,tf,darknet,caffe等。而pytorch也给我们提供了对应的接口，就是torch.onnx.export。下面具体到每一步。
首先，环境和依赖：onnx包，cuda和cudnn，我用的版本号分别是1.7.0,   10.1
TensorRT是NVIDIA推出的一款高效深度学习模型推理框架，其包括了深度学习推理优化器和运行时，能够让深度学习推理应用拥有低时延和高吞吐的优点。
本质上来讲，就是通过采用对模型中的部分算子进行融合、对特定尺寸的算子选用更好的实现方法，以及使用混合精度等方式，最终加速整个网络的推理速度。
在使用PyTorch训练得到网络模型后，我们希望在模型部署时通过TensorRT加速模型推理，那么可以先将PyTorch模型转为ONNX，然后再讲ONNX转为TensorRT的engine。
实现步骤
				RuntimeError: [enforce fail at inline_container.cc:145] . PytorchStreamReader failed reading zip arc
					13585