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
769
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
