活动地址：毕业季·进击的技术er
Pytorch导出onnx模型，C++转化为TensorRT并实现推理
本篇为学习笔记，与参考文中有出入的地方，用黄色标记出来。
主要参考：
1. Pytorch导出onnx模型，C++转化为TensorRT并实现推理过程
2. onnxruntime安装与使用（附实践中发现的一些问题）
3. TensorRT_Test
一. Pytorch导出onnx模型
-
新建一个export_onnx.py文件，全部内容如下。ResNet50_wSoftmax自定义模型，是在官方原有的基础上，添加了softmax操作。
-
将一些必要后处理添加到模型中一起导出，这样做有两个优点：
1)可以直接得到端到端的 onnx/tensorrt 模型，不必在外面再做后处理操作
2)之后我们会将 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] # 此处增加一维, 从[3,224,224] 到 [1, 3,224,224]
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是动态可变的
)
- 执行脚本：
python export_onnx.py
注意：操作之前需要准备一张图片，名为test.jpg，随便从网上下载一个即可，可以看到如下的类似的输出结果：
Downloading: "https://download.pytorch.org/models/resnet50-0676ba61.pth" to /home/cui/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 97.8M/97.8M [07:36<00:00, 225kB/s]
test_image: test.jpg, max_idx: 285, max_logit: 0.5382498502731323
二. onnxruntime推理测试
我们将刚刚得到的resnet50_wSoftmax.onnx ，用onnxruntime来进行推理测试，检测结果是否相同。
此处新建一个名为onnxruntime_test.py的文件，全部内容如下。此文件中复用了export_onnx.py中的preprocessing预处理操作，但是原函数输出结果为torch，需要将其转化为numpy类型。
# onnxruntime_test.py
import onnxruntime as ort #若缺少，使用pip install onnxruntime-gpu -i https://pypi.tuna.tsinghua.edu.cn/simple高速安装
import numpy as np
import cv2
from export_onnx import preprocessing
image_path = 'test.jpg'
ort_session = ort.InferenceSession("resnet50_wSoftmax.onnx", providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']) # 创建一个推理session
img = cv2.imread(image_path)
input_img = preprocessing(img)[None]
input_img = input_img.numpy() ## 新添加
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]}")
执行：
python onnxruntime_test.py
结果如下，基本与pytorch模型推理的一致，证明转换没问题。
test_image: test.jpg, max_idx: 285, probability: 0.5382504463195801
注意：
1）此文件ort.InferenceSession中参数与参考文献略有变动，直接执行原网站的脚本会出现如下错误。
解决参考：[ValueError: This ORT build has ‘TensorrtExecutionProvider‘, ‘CUDAExecutionProvider‘, ‘CPUExecutionP
扫描二维码关注公众号，回复：
15915797 查看本文章
ValueError: This ORT build has ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'] enabled. Since ORT 1.9, you are required to explicitly set the providers parameter when instantiating InferenceSession. For example, onnxruntime.InferenceSession(..., providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider'], ...)
2）原文直接运行推理会出现RuntimeError: Input must be a list of dictionaries or a single numpy array for input ‘image’
解决参考：onnx推理时报错RuntimeError: Input must be a list of dictionaries or a single numpy array for input ‘image
3）原文中导出的onnx模型是int64的weights，TensorRT不支持64位，但是不影响后续的使用。目前只支持32位，下一步的转换时，会报出相关类型警告。
Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32
三. onnx模型转换为tensorrt模型
1. 环境安装
CUDA10.2+Cudnn8.4.1+TensorRT7.1.3.4+Onnx-tensorrt7.1，详细的安装可以参考环境配置安装的文章，此处不多展开，例如：
- linux安装cuDNN8.4
- linux安装CUDA+cuDNN
- Ubuntu18配置CUDA10.2 cudnn8.0.1 TensorRT7.1.3.4
- Onnx-tensorrt7.1版本
- Onnx-tensorrt安装可能遇到的问题
注意：版本之间要对应，否则会出现莫名的错误。全部安装完成后，使用下面的命令进行版本检查。
nvcc -V # 查看cuda版本，10.2
cat /usr/local/cuda-10.2/include/cudnn_version.h | grep CUDNN_MAJOR -A 2 # 查看cudnn8.4版本，与之前版本稍有不同，可参考第一个链接安装cudnn
find / -name NvInferVersion.h # 查看tensorRT版本号
2. 文件夹建立
新建一个build_model.cc的C++文件，主要包含头文件，logger类，build_model函数等。
- logger类：来打印构建 tensorrt 模型过程中的一些错误或警告。按照指定的严重性程度 (severity)，来打印信息。
- build_model函数：利用onnx构建，导出engine模型（等效于trt模型，都是序列化的二进制文件）
2.1 完整文件内容如下：
// tensorrt相关
#include <NvInfer.h>
#include <NvInferRuntime.h>
#include <NvInferRuntimeCommon.h> //新添加，否则找不到nvinfer1::AsciiChar
// onnx解析器相关
#include <NvOnnxParser.h> // 与原文不同，onnx-tensorrt build后，sudo make install
// cuda_runtime相关
#include <cuda_runtime.h>
// 常用头文件
#include <math.h>
#include <stdio.h>
#include <unistd.h>
#include <chrono>
#include <fstream>
#include <functional>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include <dirent.h>
// opencv
#include <opencv2/opencv.hpp>
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, const char* 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);
}
}
};
bool build_model() {
if (access("resnet50.engine", 0) == 0) {
printf("resnet50.engine 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("../resnet50_wSoftmax.onnx", 1)) {
printf("Failed to parse resnet50_wSoftmax.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("resnet50.engine", "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;
}
int main() {
if (!build_model()) {
printf("Couldn't build engine!\n");
}
return 0;
}
2.2 Cmakelists.txt文件
完整工程参考：TensorRT_Test
2.3 编译运行
mkdir build && cd build
cmake ..
make -j8
./build_model
2.4 结果
出现如下情况，且生成了rsenet50.engine文件，即为成功转换。
warning: [TRT]/home/cui/workspace/tools/onnx-tensorrt/onnx2trt_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
Workspace Size = 256.00 MB
Build Done.
2.5 可能遇到的问题
Q1：error: ‘nvinfer1::AsciiChar’ has not been declared
解决：tensorRT版本问题，将其改为const char*。
Q2: error: Could not open file resnet50_wSoftmax.onnx
解决：修改build_model.cc文件中78行，resnet50_wSoftmax.onnx 模型文件的对应路径
四. TensorRT模型推理测试
目的：验证与之前的Pytorch模型和onnx模型推理结果是否一致。
- 完整文件内容，请参考工程中的model.infer.cc文件，与原文中稍有改动。增加了checkRuntime实现：
#ifndef checkRuntime
#define checkRuntime(callstr)\
{\
cudaError_t error_code = callstr;\
if (error_code != cudaSuccess) {\
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__ << std::endl;\
assert(0);\
}\
}
#endif // checkRuntime
工程地址：TensorRT_Test
- 编译执行
mkdir build && cd build
cmake ..
make -j8
./model_infer
- 结果
test_image: ../test.jpg, max_idx: 285, probability: 0.538250(learning3d)
- 可能遇到的问题
Q1: error while loading shared libraries: libcudnn.so.8: cannot open shared object file: No such file or directory
解决：未找到cuda安装环境，可以进入conda配置好的某个环境下，在执行相应的操作。或者参看CUDA安装环境，将其配置到环境变量中。