6.2.tensorRT高级(1)-第一个完整的分类器程序
目录
-
- 前言
- 1. CNN分类器
- 2. 补充知识
-
- 2.1 知识点
- 2.2 智能指针封装
- 总结
前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 高级-第一个完整的分类器程序
课程大纲可看下面的思维导图
1. CNN分类器
这节课我们学习 tensorRT 高级篇案例——CNN分类器
先来看下代码，gen_onnx.py 是推理一个 resnet18 的 ImageNet 分类器模型，打印其推理结果，并导出 onnx 模型，其具体内容如下：
import torch
import torchvision
import cv2
import numpy as np
class Classifier(torch.nn.Module):
def __init__(self):
super().__init__()
#使用torchvision自带的与训练模型, 更多模型请参考:https://tensorvision.readthedocs.io/en/master/
self.backbone = torchvision.models.resnet18(pretrained=True)
def forward(self, x):
feature = self.backbone(x)
probability = torch.softmax(feature, dim=1)
return probability
# 对每个通道进行归一化有助于模型的训练
imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std = [0.229, 0.224, 0.225]
image = cv2.imread("workspace/dog.jpg")
image = cv2.resize(image, (224, 224)) # resize
image = image[..., ::-1] # BGR -> RGB
image = image / 255.0
image = (image - imagenet_mean) / imagenet_std # normalize
image = image.astype(np.float32) # float64 -> float32
image = image.transpose(2, 0, 1) # HWC -> CHW
image = np.ascontiguousarray(image) # contiguous array memory
image = image[None, ...] # CHW -> 1CHW
image = torch.from_numpy(image) # numpy -> torch
model = Classifier().eval()
with torch.no_grad():
probability = model(image)
predict_class = probability.argmax(dim=1).item()
confidence = probability[0, predict_class]
labels = open("workspace/labels.imagenet.txt").readlines()
labels = [item.strip() for item in labels]
print(f"Predict: {predict_class}, {confidence}, {labels[predict_class]}")
dummy = torch.zeros(1, 3, 224, 224)
torch.onnx.export(
model, (dummy,), "workspace/classifier.onnx",
input_names=["image"],
output_names=["prob"],
dynamic_axes={"image": {0: "batch"}, "prob": {0: "batch"}},
opset_version=11
)
我们使用的 resnet18 来自于 torchvision.models，它的输出值并不是一个概率值，需要再加上一个 softmax 函数才行，因此我们自己实现了一个类顺便把 softmax 操作也添加进去，减少后处理的复杂度，导出的 onnx 包含后处理直接输出概率值
接下来就是对输入图像进行预处理操作，得到一个 image 的 tensor 放到网络中推理，前面有提到图像预处理的操作可以用 warpAffine 来实现，一个核函数解决所有操作，使得其性能足够高，效果足够好
值的注意的是，在 onnx 导出时有两个细节，一个是输入和输出的动态维度都只指定 batch 动态，宽高不动态，另一个是 optset 的版本指定为 11
运行效果如下：
我们接下来到 C++ 里面看下 tensorRT 是怎么处理 CNN 分类器的推理过程的，代码如下：
// tensorRT include
// 编译用的头文件
#include 由于是通过 onnx 模型构建网络，因此要包含 NvOnnxParser，依然使用的是我们自定义的源代码，而不是通过 libnvonnxparser.so 库文件来解析
先来看模型构建 build_model 与之前不同的是对于 tensorRT API 的指针我们又套了一层，通过 make_nvshared() 函数来实现创建 tensorRT 对象的智能指针，它可以自动管理其所指向的对象的生命周期，实现内存的自动释放，这是 make_nvshared 完成的工作，其它流程与之前都一样
再来看下 inference，之前的输入都是手动给的，这次需要给定一张图像，首先需要根据输入图像的大小分配 Host 和 Device 的大小，接下来需要读取图片完成与 pytorch 一致的预处理过程
预处理代码如下：
// BGR -> RGB
// (pixel / 255.0 - mean) / std
// to tensor -> BGRBGRBGR -> BBBGGGRRR
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){
// 注意这里的顺序rgb调换了
*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];
}
它具体是如何实现的呢？我们可以看下图：
opencv 的图像格式如左边的图所示，BGRBGRBGR… 这样连续的，而我们实际想要的是右边的效果就是一个通道一个通道的排列起来
那你可能有所困惑，右图中 phost_b 的地址中为什么填充的都是 R 像素值，而反观 phost_r 的地址中填充的都是 B 像素值？是不是画错了？那是没有画错的，在程序中它其实是有故意把 R 像素值填充到 b 通道的，而把 B 像素值填充到 r 通道的，这样做的目的当然就是 bgr2rgb 这个预处理操作啦
我们需要转换的首地址当然就是 input_data_host，先取得 b 通道的首地址 phost_b，然后通过图像面积的偏移量分别获取 phost_g、phost_r，然后将 opencv 对应的 b，g，r 分别塞到对应的地方就行，在这个过程中由于是直接对像素进行操作，因此非常容易实现诸如 BRG2RGB、/255.0、减均值除标准差等各种预处理操作，简单且高效，非常值得学习。
input_data_host 的处理完成后可以将预处理好的图像数据拷贝到 input_data_device 上面，input_data_device 准备好了之后，可以开始准备 output_device 了，根据模型预测的 1000 个类别数分配一个 output_data_device 的空间，分配好之后，可以指定下推理的 batch 数目，然后指定下 bingding 的输入和输出，把它塞到 enqueueV2 中就可以去进行推理了，推理完成之后将推理结果从 output_data_device 拷贝到 output_data_host 上面，接下来做一个流同步，等待所有操作结束就好了。
整个流程图如下所示：
从上面流程可以看到整个过程非常清晰，而典型的 CUDA 核函数程序的执行流程和上述差不多，无非是包含以下几个步骤：
1. 分配 host 内存，并进行数据的初始化
2. 分配 device 内存，并从 host 将数据拷贝到 device 上
3. 调用 CUDA 的核函数在 device 上完成指定的计算
4. 将 device 上的运算结果拷贝到 host 上
5. 释放 device 和 host 上分配的内存
最终的 output_data_host 就是我们想要的概率值，通过 1000 个类别的概率值可以拿到最大概率值对应的标签，然后将其打印出来，输出结果如下：
可以看到和 python 结果一致，说明我们整个中间过程没有任何问题，这就是分类器这个案例，上述代码的重点是预处理，对预处理的理解有利于后续实现 CUDA 核上的并发方式
2. 补充知识
对于上述代码的分析，你需要知道：
1. 模型里面没有 softmax 操作，在这里采用包裹一层加上 softmax 节点后再导出模型，这样使得后处理得到的直接就是概率值，避免后处理上再做 softmax
2. 在 c++ 代码中，则充分采用指针偏移的方式，提升 cpu 上预处理的效率
3. 对于 bgr->rgb 也避免使用 cvtColor 实现，而是简单的改变赋值时的索引，提升效率
2.1 知识点
关于分类器案例的知识点：(from 杜老师)
1. 使用智能指针，对 tensorrt 返回值做封装，使得内存安全不会泄露
template<typename _T>
shared_ptr<_T> make_nvshared(_T* ptr){
return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}
// [](_T* p){p->destroy();} 这里用 lambda 表达式的形式来表示 destroy 的方式
// 因为它常常需要 destroy 进行释放
2. 使用 cudaMallocHost 对输入的 host 进行分配，使得主机内存复制到设备效率更高
3. 注意推理时的预处理，指定了 rgb 与 bgr 对调
4. 如果需要多个图像推理，需要：
- 在编译时，指定 maxbatchsize 为多个图
- 在推理时，指定输入的 bindings shape 的 batch 维度为使用的图像数，要求小于等于 maxbatchsize
- 在收取结果的时候，tensor 的 shape 是 input 指定的 batch 大小，按照 batch 处理即可
2.2 智能指针封装
我们来简单聊下智能指针的封装，代码如下：
template<typename _T>
shared_ptr<_T> make_nvshared(_T* ptr){
return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});
}
上述代码实现了一个模板函数，用于创建一个包含 NVIDIA TensorRT 对象的智能指针（例如 nvinfer1::ICudaEngine, nvinfer1::IExecutionContext 等）。这个智能指针的特点是，当它超出作用域或者没有任何引用时，它会自动删除其所指向的对象，从而避免内存泄漏。这是通过 C++ 的 std::shared_ptr 来实现的。
我们来具体分析下这段代码：
1. template：这是一个模板声明，表示这个函数可以接受任何类型的参数。_T 是一个占位符，表示函数接受的函数类型。_T 只是一种常见的命名约定，用于表示模板参数，你甚至可以起名叫做 ABC，这都没关系
2. shared_ptr<_T> make_nvshared(_T* ptr)：这是函数声明，函数名为 make_nvshared，它接受一个指向 _T 类型的指针 ptr，并返回一个模板类 std::shared_ptr<_T>，它是一个智能指针，可以自动管理其所指向的对象的生命周期
3. return shared_ptr<_T>(ptr, [](_T* p){p->destroy();});：这是函数的主体部分，它创建并返回一个智能指针实例对象 std::shared_ptr<_T>。std::shared_ptr<_T> 的构造函数接受两个参数：第一个参数是它应该管理的对象的指针，第二个参数是一个删除器。删除器是一个函数或者可调用对象，当智能指针不再需要其所指向的对象时，它会调用删除器来删除这个对象。在这里，删除器是一个 lambda 表达式 [](_T p){p->destroy();}*，当智能指针不再需要其所指向的对象时，它会调用这个对象的 destroy 成员函数来删除它。
所以，这段代码的作用是创建一个智能指针，这个智能指针会自动删除其所指向的对象，从而避免内存泄漏。这是通过使用 std::shared_ptr 的删除器来实现的，删除器是一个 lambda 表达式，它调用对象的 destroy 成员函数来删除对象。
简单来说，上述代码实现了一个函数，接收一个指针，返回对应类型的智能指针实例对象。
有关共享智能指针相关的知识看参考苏老师的 shared_ptr
总结
本次课程学习了 tensorRT 高级篇的第一个案例-CNN分类器，在高级篇里面就要有高级篇的做法，因此我们使用智能指针对 tensorrt 对象的返回值进行了封装，防止内存的泄露。同时我们在 CPU 上也实现了一个高性能的预处理方法，直接对像素进行操作，可以同时把几个预处理操作都做了，对 CPU 的图像预处理操作的理解有利于我们后续实现 warpAffine CUDA 核函数。