Torch 模型 model =＞ .onnx =＞ .trt 及利用 TensorTR 在 C++ 下的模型部署教程

一、模型训练环境搭建和模型训练

模型训练环境搭建主要牵扯 Nvidia driver、Cuda、Cudnn、Anaconda、Torch 的安装，相关安装教程可以参考【StarCoder 微调《个人编程助手: 训练你自己的编码助手》】中 5.1 之前的章节。

模型训练的相关知识可以参考 Torch的编程方式示例以及必备知识，可以直接使用第一部分的手写字体分类模型作为当前教程的示例模型，只需要在脚本最后加入如下代码便可将模型保存为 .pth 模型。注意在上载 model 的代码里依然还是要定义一遍模型，或者 import 之前的模型定义 py 文件，否则报错。

torch.save(model, './model.pth')
# model = torch.load('./model.pth')
# 虽然保存了模型结构和模型参数，但是在load之前还是要定义一遍model才行，或者直接import model

二、model 转换 .onnx 并在 GPU 上利用 python 代码推理

Torch model 模型通过如下代码可以转换为 .onnx 模型。参考链接。

torch.onnx.export(model, input, onnx_path, verbose=True, opset_version=12, input_names=input_names, output_names=output_names) 

# model 不是 .pth 路径，而是读取内存中的模型，需要先引入模型定义的类，再 load .pth 文件得到 model
# input 是 model 的 forward 函数的输入，如果输入只有一个图片，那 input 可以是一个张量或者字典，但如果输入有多个，则 imput 只能是字典，键名与 forward 形参名称对应
# onnx_path 是输出 onnx 的路径
# verbose 是转换过程中是否打印详情
# input_names 定义 onnx 中输入的名称列表
# output_names 定义 onnx 中输出的名称列表

如果想要在 GPU 上利用 python 代码进行 onnx 模型推理，那首先需要安装 onnxruntime-gpu，安装时需要版本匹配，需要与 Cuda 和 Cudnn 版本匹配，官方版本匹配说明参考链接。安装命令如下：

pip install onnxruntime-gpu==1.18.1 -i https://pypi.tuna.tsinghua.edu.cn/simple

在完成 onnxruntime-gpu 的安装后，可以通过如下代码完成模型推理。但推理性能一般，其在 Cuda 下的推理时间似乎大于在 Torch 下的推理时间。参考链接。

providers = ['CUDAExecutionProvider']
m = rt.InferenceSession(onnx_path, providers=providers)

# 推理 onnx 模型
# output_names 为导出 onnx 时定义的输出名称列表，例如 output_names = ['hm', 'vaf', 'haf','curb_hm', 'curb_vaf', 'curb_haf']
# {"input": image} 为导出 onnx 时定义的输入名称列表
onnx_pred = m.run(output_names, {"input": image})

三、TensorRT 安装

TensorRT 存在 Python 版本和 C++ 版本。

Python 版本 TensorRT 的安装很简单。但一般不太会用 Python 版本 TensorRT。

// TensorRT 8.5 以上的版本用以下命令
pip install tensorrt 

// 测试一下
python
>>> import tensorrt
>>> print(tensorrt.__version__)
>>> assert tensorrt.Builder(tensorrt.Logger())

C++ 版本 TensorRT 的安装分为.deb 和 . tar 的方式，一般选用 .tar 的方式。

首先去官网下载 TensorRT，下载时版本即需要与 Cuda 版本匹配，也需要与 Cudnn 匹配。官网中只写了 TensorRT 与 Cuda 如何匹配，所以与 Cudnn 的匹配需要自己试出来，当前成功安装的版本是 Cuda12.1 + Cudnn 9.7.1 + TensorRT 10.0.0.6，测试发现 TensorRT 8.* 要求 Cudnn 8.*。

tar -xzvf TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
cp -r TensorRT-8.4.3.1 /usr/local/include/TensorRT-8.4.3.1
// 路径无所谓，只要知道在哪里就可以，下面 bashrc 中写的就是这个路径，
// cp 完后 home 目录下可删除，或者直接不 cp 而是直接使用 home 下这个路径也可以

sudo vim .bashrc

// 添加如下内容，注意版本号和路径的对应
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/include/TensorRT-8.4.3.1/lib
alias trtexec="/usr/local/include/TensorRT-8.4.3.1/bin/trtexec"

# 如果是 deb 安装的，则可能是：
alias trtexec="/usr/src/tensorrt/bin/trtexec"
# or 
alias trtexec="/usr/local/tensorrt/bin/trtexec"


source ~/.bashrc

安装完 TensorRT 后，在 vscode 里直接 #include <NvInfer.h> 会报错找不到，解决办法是在 vscode 设置 (.vscode/c_cpp_properties.json) 里添加如下 TensorRT 和 cuda 两个路径，vscode 添加 include 路径的办法参考链接。

{
    "configurations": [
        {
            "name": "Linux",
            "includePath": [
                "${workspaceFolder}/**",
                "/usr/local/include/TensorRT-10.0.0.6/include",
                "/usr/local/cuda-12.1/include"
            ],
            "defines": [],
            "compilerPath": "/usr/bin/gcc",
            "cStandard": "c17",
            "cppStandard": "gnu++17",
            "intelliSenseMode": "linux-gcc-x64"
        }
    ],
    "version": 4
}

四、.onnx 模型转 .trt 模型

.onnx 模型转 .trt 模型有两种方法，第一种方法是使用现成的 trtexec 命令，第二种方法是使用自己写的 C++ 代码，第二种方式可以参考链接 1 链接 2 链接3。这里主要使用第一种方法，关于第一种方法的详细介绍可以参考链接。大部分情况只需要用到下面这两句命令。

trtexec --onnx=2Dmodel.onnx --saveEngine=2Dmodel.trt

// 转换为 fp16 模型
trtexec --onnx=2Dmodel.onnx --saveEngine=2Dmodel.trt --fp16

五、.trt 模型在 C++ 下进行推理

首先，介绍两个概念：

CUDA 流：可以将一个流看做是 GPU 上的一个任务，不同任务可以并行执行。利用三个流，同一个流上的任务顺序执行，不同流上的任务可以同时执行，从而实现并发操作。

运行时库：就是在运行时提供基础支持的库。例如：C/C++ 运行时库，就是 C/C++ 程序在运行时依赖的基础库。

然后，说明一下官方资料：

Nvidia TensorRT 官方文档 (最新版本历史版本)、Quick Start Guide (最新版本历史版本可以从官方文档进入)、API 说明 (最新版本历史版本可以从官方文档进入)、版本变化 API 变化说明和代码 (最新版本历史版本)。

这里存在一个情况，TensorRT 从 8.x 版本变为 10.x 版本后，API 发生了一些变化，例如：模型推理函数从 enqueueV2 变为 enqueueV3，且形参列表也变了；还有新版本似乎不需要再主动调用 destroy() 函数进行对象销毁。但官方文档、Quick Start Guide 和代码等，更新并不及时，到了最新版才完成更新，所以如果你实际安装的是 TensorRT 10.0.0.6，然后你参考对应版本的官方资料，会无法运行。

最后，说明一下第三方资料：

当前我搜索到的所有第三方 CSDN 和知乎资料全是针对 8.x 版本的。但也很值得参考

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

torch 转 ONNX 模型转 TensorRT C++ 推理

TensorRT 的使用流程 c++ 和 python

下面是我针对包含三个 head (语义分割、关键点检测、旋转框 FCOS 检测) 的 trt 模型的推理代码，针对版本是 Cuda12.1 + Cudnn 9.7.1 + TensorRT 10.0.0.6。

FP16 的推理与 FP32 的推理代码没区别，唯一区别就是上一步转 .trt 时多了个配置参数，所以在模型推理时，输入依然是 FP32，进入模型后，模型会自动转换为 FP16，输出也一样，模型会自动转换为 FP32。实测在 24G A10 下，100 次推理，512 × 256 尺寸输入图片，FP32 推理耗时 394ms，FP32 推理耗时 253ms。

IPMdet.h

#pragma once
// 系统头文件
#include <string>
#include <vector>
#include <fstream>
#include <memory>
// TensorRT 相关头文件
#include <NvInfer.h>
#include <NvInferRuntime.h>
// CUDA 相关头文件
#include <cuda_runtime.h>
// 测试用头文件
#include <iostream>
using namespace std;


// 用于打印 CUDA 报错 - 宏定义
#define checkRuntime(op) __check_cuda_runtime((op), #op, __FILE__, __LINE__)


// 用于打印 TensorRT 报错，准备一个 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);
		}
	}
};

// 定义一个 share_ptr 的智能指针
// 模板的定义必须放在头文件中，因为模板的实例化需要在编译时进行
template<typename _T>
shared_ptr<_T> make_nvshared(_T *ptr) {
	return shared_ptr<_T>(ptr);
}

// IPM 图检测类
class IPMdet {
	public:
		IPMdet() = default;
		IPMdet(const string& file);
		~IPMdet();
		// 模型推理
		void infer_trtmodel(float* pimage);

	private:
		// 读取 .trt 文件
		vector<unsigned char> load_file(const string& file);
		// 预处理输入图片
		void preprocess_image(float* pimage);

	public:
		// 输出结果

    private:
		// TensorRT 推理用的工具
		vector<unsigned char> _engine_data;                           // 记录 .trt 模型的二进制序列化格式数据
		TRTLogger logger;                                             // 打印 TensorRT 的错误信息
		shared_ptr<nvinfer1::IRuntime> _runtime = nullptr;            // 运行时，即推理引擎的支持库和函数等
		shared_ptr<nvinfer1::ICudaEngine> _engine = nullptr;          // 推理引擎，包含反序列化的 .trt 模型数据
		shared_ptr<nvinfer1::IExecutionContext> _context = nullptr;   // 上下文执行器，用于做模型推理
		cudaStream_t _stream = nullptr;

		// 定义模型输入输出尺寸
		int input_batch = 1;
		int input_channel = 3;
		int input_height = 512;
		int input_width = 256;
		int output_height1 = input_height / 4;
		int output_width1 = input_width / 4;
		int output_height2 = input_height / 8 ;
		int output_width2 = input_width / 8;
	
		// 准备好 **_host 和 **_device，分别表示内存中的数据指针和显存中的数据指针
		// input 数据
		int input_numel = input_batch * input_channel * input_height * input_width;
		float* input_data_host = nullptr;
		float* input_data_device = nullptr;

		// output 数据 - keypoints
		int keypoints_numel = input_batch * 1 * output_height1 * output_width1;
		void* output_data_keypoints_host = nullptr;
		void* output_data_keypoints_device = nullptr;

		// output 数据 - classification
		int classification_numel = input_batch * 1 * output_height2 * output_width2;
		void* output_data_classification_host = nullptr;
		void* output_data_classification_device = nullptr;

		// output 数据 - bbox_OBB
		int bbox_OBB_numel = input_batch * 4 * output_height2 * output_width2;
		void* output_data_bbox_OBB_host = nullptr;
		void* output_data_bbox_OBB_device = nullptr;

		// output 数据 - centerness
		int centerness_numel = input_batch * 1 * output_height2 * output_width2;
		void* output_data_centerness_host = nullptr;
		void* output_data_centerness_device = nullptr;

		// output 数据 - bbox_wh
		int bbox_wh_numel = input_batch * 2 * output_height2 * output_width2;
		void* output_data_bbox_wh_host = nullptr;
		void* output_data_bbox_wh_device = nullptr;

		// output 数据 - segmentation
		int segmentation_numel = input_batch * 8 * input_height * input_width;
		void* output_data_segmentation_host = nullptr;
		void* output_data_segmentation_device = nullptr;
};

IPMdet.cpp

#include "IPMdet.h"


// 打印 Cuda 报错
// 函数定义不应放在头文件里，否则当头文件被多个 cpp 调用时，会出现重复定义的问题
bool __check_cuda_runtime(cudaError_t code, const char* op, const char* file, int line) {
	if (code != cudaSuccess) {
		const char* err_name = cudaGetErrorName(code);
		const char* err_message = cudaGetErrorString(code);
		printf("runtime error %s: %d  %s failed.\n  code = %s, message = %s", file, line, op, err_name, err_message);
		return false;
	}
	return true;
}

// 构造函数
IPMdet::IPMdet(const string& file){
    // 1. TensorRT 相关操作
    ///
    // 1.1 读取 .trt 文件
    _engine_data = load_file(file);

    // 1.2 创建运行时，需要日志记录器
    _runtime = make_nvshared(nvinfer1::createInferRuntime(logger));

    // 1.3 创建推理引擎，需要运行时和序列化 trt 文件，包含反序列化的 .trt 模型数据
    _engine = make_nvshared(_runtime->deserializeCudaEngine(_engine_data.data(), _engine_data.size()));
    if (_engine == nullptr) {
        printf("Deserialize cuda engine failed.\n");
        return;
    }

    // 1.4 创建上下文执行器，需要推理引擎
    _context = make_nvshared(_engine->createExecutionContext());

    // 打印 .trt 模型的输入输出张量的名称和维度，这里与 onnx 中的名称和维度一致
    // for (int i=0, e=_engine->getNbIOTensors(); i<e; i++){
    //     auto const name = _engine->getIOTensorName(i);
    //     auto const size = _engine->getTensorShape(name);
    //     cout << "Tensor Name: " << name << endl;
    //     for (int i = 0; i < size.nbDims; ++i) {
    //         std::cout << "Dimension " << i << ": " << size.d[i] << std::endl;
    //     }
    //     cout << endl;
    // }

    // 2. CUDA 相关操作
    ///
    // 2.1 创建 CUDA 流，CUDA 流类似于线程，每个任务都必须有一个 CUDA 流，不同的 CUDA 流可以在 GPU 中并行执行任务
	checkRuntime(cudaStreamCreate(&_stream));

    // 2.2 申请 CPU 内存和 GPU 内存，准备好 **_host 和 **_device，分别表示内存中的数据指针和显存中的数据指针
	// input 数据
    checkRuntime(cudaMallocHost(&input_data_host, input_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&input_data_device, input_numel * sizeof(float)));

    // output 数据 - keypoints
    checkRuntime(cudaMallocHost(&output_data_keypoints_host, keypoints_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_keypoints_device, keypoints_numel * sizeof(float)));

    // output 数据 - classification
    checkRuntime(cudaMallocHost(&output_data_classification_host, classification_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_classification_device, classification_numel * sizeof(float)));

    // output 数据 - bbox_OBB
    checkRuntime(cudaMallocHost(&output_data_bbox_OBB_host, bbox_OBB_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_bbox_OBB_device, bbox_OBB_numel * sizeof(float)));

    // output 数据 - centerness
    checkRuntime(cudaMallocHost(&output_data_centerness_host, centerness_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_centerness_device, centerness_numel * sizeof(float)));

    // output 数据 - bbox_wh
    checkRuntime(cudaMallocHost(&output_data_bbox_wh_host, bbox_wh_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_bbox_wh_device, bbox_wh_numel * sizeof(float)));

    // output 数据 - segmentation
    checkRuntime(cudaMallocHost(&output_data_segmentation_host, segmentation_numel * sizeof(float)));
	checkRuntime(cudaMalloc(&output_data_segmentation_device, segmentation_numel * sizeof(float)));

    // 3. TensorRT 内存绑定
    ///
    _context->setTensorAddress("image", input_data_device);
    _context->setTensorAddress("keypoints", output_data_keypoints_device);
    _context->setTensorAddress("classification", output_data_classification_device);
    _context->setTensorAddress("bbox_OBB", output_data_bbox_OBB_device);
    _context->setTensorAddress("centerness", output_data_centerness_device);
    _context->setTensorAddress("bbox_wh", output_data_bbox_wh_device);
    _context->setTensorAddress("segmentation", output_data_segmentation_device);

    cout << "init OK !" << endl;
}

// 析构函数
IPMdet::~IPMdet(){
    // 释放 CPU 内存，看 TensorRT 10 以上版本，似乎不再需要主动释放 CPU 内存
    checkRuntime(cudaFreeHost(input_data_host));
    checkRuntime(cudaFreeHost(output_data_keypoints_host));
    checkRuntime(cudaFreeHost(output_data_classification_host));
    checkRuntime(cudaFreeHost(output_data_bbox_OBB_host));
    checkRuntime(cudaFreeHost(output_data_centerness_host));
    checkRuntime(cudaFreeHost(output_data_bbox_wh_host));
    checkRuntime(cudaFreeHost(output_data_segmentation_host));

    // 释放 GPU 内存
	checkRuntime(cudaFree(input_data_device));
	checkRuntime(cudaFree(output_data_keypoints_device));
    checkRuntime(cudaFree(output_data_classification_device));
    checkRuntime(cudaFree(output_data_bbox_OBB_device));
    checkRuntime(cudaFree(output_data_centerness_device));
    checkRuntime(cudaFree(output_data_bbox_wh_device));
    checkRuntime(cudaFree(output_data_segmentation_device));

    // 释放 CUDA 流，看 TensorRT 10 以上版本，似乎不再需要主动释放 CUDA 流
    checkRuntime(cudaStreamDestroy(_stream));

    cout << "Detroy OK !" << endl;
}

// 读取 .trt 文件
vector<unsigned char> IPMdet::load_file(const string& file) {  // 返回结果为无符号字符的vector，其数据存储是连成片的
    ifstream in(file, ios::in | ios::binary);          // 定义一个数据读取对象，以二进制读取数据
    if (!in.is_open()) return {};                      // 如果没有可读数据则返回空

    in.seekg(0, ios::end);                             // seekg函数作用是将指针指向文件终止处
    size_t length = in.tellg();                        // tellg函数作用是返回指针当前位置，此时即为数据长度

    vector<uint8_t> data;                              // 定义一个vector用于存储读取数据，仅仅是类头，其数据存储区还是char型data指针
    if (length > 0) {
         in.seekg(0, ios::beg);                        // seekg函数作用是将指针指向文件起始处
         data.resize(length);                          // 为vector申请长度为length的数据存储区，默认全部填充 0
         in.read((char*)&data[0], length);             // 为vector的数据存储区读取长度为length的数据
    }
    in.close();                                        // 关闭数据流
    return data;
}

// 预处理输入图片
void IPMdet::preprocess_image(float* pimage){

  float mean[] = {128, 128, 128};
  float std[] = {128, 128, 128};

  int image_area = 512 * 256;
  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] - mean[0]) / std[0];
       *phost_g++ = (pimage[1] - mean[1]) / std[1];
       *phost_b++ = (pimage[2] - mean[2]) / std[2];
   }
}

void IPMdet::infer_trtmodel(float* pimage){
    // 预处理输入图片
    preprocess_image(pimage);
    // 将输入图片从 CPU 内存拷贝至 GPU 内存
    checkRuntime(cudaMemcpyAsync(input_data_device, input_data_host, input_numel *sizeof(float), cudaMemcpyHostToDevice, _stream));
    // 模型推理
    bool success = _context->enqueueV3(_stream);
    // 将输出结果从 GPU 内存拷贝至 CPU 内存
    checkRuntime(cudaMemcpyAsync(output_data_keypoints_host, output_data_keypoints_device, sizeof(output_data_keypoints_host), cudaMemcpyDeviceToHost, _stream));
    checkRuntime(cudaMemcpyAsync(output_data_classification_host, output_data_classification_device, sizeof(output_data_classification_host), cudaMemcpyDeviceToHost, _stream));
    checkRuntime(cudaMemcpyAsync(output_data_bbox_OBB_host, output_data_bbox_OBB_device, sizeof(output_data_bbox_OBB_host), cudaMemcpyDeviceToHost, _stream));
    checkRuntime(cudaMemcpyAsync(output_data_centerness_host, output_data_centerness_device, sizeof(output_data_centerness_host), cudaMemcpyDeviceToHost, _stream));
    checkRuntime(cudaMemcpyAsync(output_data_bbox_wh_host, output_data_bbox_wh_device, sizeof(output_data_bbox_wh_host), cudaMemcpyDeviceToHost, _stream));
    checkRuntime(cudaMemcpyAsync(output_data_segmentation_host, output_data_segmentation_device, sizeof(output_data_segmentation_host), cudaMemcpyDeviceToHost, _stream));
    // 等待直到 _stream 流的工作完成
    checkRuntime(cudaStreamSynchronize(_stream));
}

main.cpp

#include <iostream>
#include <string>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "IPMdet.h"
using namespace std;

int main() {
    cv::Mat image = cv::imread("/mnt/sdb/ipm_sample_datas/detection/train/apa_det_train/ADAS_20200113-132415_216_1212__00006732_28x14_03_ipm.jpg");
    cout << image.size << endl;

    cv::Mat input;
    image.convertTo(input, CV_32FC3);

    string onnx_path = "/mnt/sdb/HAT-feature-ipm-multitask-release/output_model/model_best_22.trt";
    IPMdet ipmdet = IPMdet(onnx_path);
    auto start0 = chrono::high_resolution_clock::now();
    ipmdet.infer_trtmodel((float*)input.data);
    
    for(int i=0; i<100; i++){
        ipmdet.infer_trtmodel((float*)input.data);
    }
    auto end0 = chrono::high_resolution_clock::now();

    string onnx_path_fp16 = "/mnt/sdb/HAT-feature-ipm-multitask-release/output_model/model_best_22_fp16.trt";
    IPMdet ipmdet_fp16 = IPMdet(onnx_path_fp16);
    ipmdet_fp16.infer_trtmodel((float*)input.data);

    auto start1 = chrono::high_resolution_clock::now();
    for(int j=0; j<100; j++){
        ipmdet_fp16.infer_trtmodel((float*)input.data);
    }
    auto end1 = chrono::high_resolution_clock::now();

    chrono::duration<double, std::milli> elapsed0 = end0 - start0;
    chrono::duration<double, std::milli> elapsed1 = end1 - start1;
    cout << "FP32 time taken: " << elapsed0.count() << " ms" << endl;
    cout << "FP16 time taken: " << elapsed1.count() << " ms" << endl << endl;
}