【CV-CUDA实战】使用Python+TensorRT+CVCUDA优化YOLOv8

原创

已于 2024-05-05 23:05:06 修改 · 2.7k 阅读

标签

#YOLO #人工智能 #python #计算机视觉

于 2024-05-05 23:03:33 首次发布

什么是CV-CUDA

NVIDIA CV-CUDA™ 是一个开源项目，用于构建云规模人工智能 (AI) 成像和计算机视觉 (CV) 应用程序。它使用图形处理单元 (GPU) 加速来帮助开发人员构建高效的预处理和后处理管道。它可以将吞吐量提高 10 倍以上，同时降低云计算成本。

环境准备

准备CV-CUDA静态库

首先我们需要访问CV-CUDA Github仓库
去下载与您系统架构相对应的静态库文件，以我为例，我需要下载：

cvcuda-dev-0.7.0_beta-cuda11-x86_64-linux.tar.xz
cvcuda-lib-0.7.0_beta-cuda11-x86_64-linux.tar.xz
cvcuda-python3.10-0.7.0_beta-cuda11-x86_64-linux.tar.xz
cvcuda-tests-0.7.0_beta-cuda11-x86_64-linux.tar.xz

请注意！下载的静态库的版本取决你的CUDA版本和目前正在使用的Python版本（这是由于PyBind11导致的）。

解压

tar -xvf cvcuda-dev-0.7.0_beta-cuda11-x86_64-linux.tar.xz
tar -xvf cvcuda-lib-0.7.0_beta-cuda11-x86_64-linux.tar.xz
tar -xvf cvcuda-python3.10-0.7.0_beta-cuda11-x86_64-linux.tar.xz
tar -xvf cvcuda-tests-0.7.0_beta-cuda11-x86_64-linux.tar.xz

添加至变量

所有解压后的静态库文件会被解压到一个名为opt的目录中，但是请注意，这个opt不是Linux根文件系统的/opt，请务必做出区分！
请在~/.bashrc文件中添加以下内容：

export LD_LIBRARY_PATH=$HOME/opt/nvidia/cvcuda0/lib/x86_64-linux-gnu

接下来使~/.bashrc生效并加载静态库

source ~/.bashrc && ldconfig

使用ldconfig命令检查静态库是否被加载：

ldconfig -p | grep nvcv

如果看到以下结果说明静态库已被加载成功：

libnvcv_types.so.0 (libc6,x86-64) => /usr/local/cuda/targets/x86_64-linux/lib/libnvcv_types.so.0

将PyBind静态库复制到env下

在您的解压目录下，应该有一个名为cvcuda.cpython-310-x86_64-linux-gnu.so的文件，请将此文件复制到以下路径：

$PYTHON_HOME/site-packages/

以我的情况为例，我所使用的是anaconda3，因此这个路径是：

/home/elin/anaconda3/envs/jetsonsim/lib/python3.10/site-packages

使用import命令来检查cvcuda是否可用：

>>> import cvcuda
>>>

如果没有报错则说明cvcuda已经被加载成功了。

算子设计

前处理算子

由于YOLOv8模型接受一个640x640，batch_size为1的RGB图像，因此我们的前处理算子可以设计成：

def preprocess(self,imageFrame:np.ndarray) -> torch.Tensor:
        # 首先将输入的Numpy矩阵转换成Pytorch张量
        imageFrame = torch.tensor(imageFrame,device="cuda",dtype=torch.uint8)
        self.imageHeight,self.imageWidth = imageFrame.shape[:2]
        # 由于OpenCV读取的图像是HWC格式，
        # 因此CVCUDA需要将图像张量按此格式进行转换
        imageTensor = cvcuda.as_tensor(imageFrame,"HWC")
        # 使用GPU上的色域转换和resize算子进行处理
        imageTensor = cvcuda.cvtcolor(imageTensor,cvcuda.ColorConversion.BGR2RGB)
        imageTensor = cvcuda.resize(imageTensor,(self.inputWidth,self.inputHeight,3))
        # 将CVCUDA张量转移成Pytorch张量
        # 由于CVCUDA张量并未实现运算符重载，所以这一步是必要的
        # 以便我们处进行二值化和后续TensorRT模型的推理
        imageData = torch.as_tensor(imageTensor.cuda(),device="cuda")
        imageData = imageData / 255.0
        # 使用维度转置将图像扩增至(1,3,640,640)
        imageData = imageData.transpose(0,2).transpose(1,2).unsqueeze(0)
        return imageData # 返回GPU张量，以避免TensorRT推理前的GPU显存拷贝

以下是前处理算子的工作示意图：
请添加图片描述

TensorRT模型加载

在TensorRT模型的加载上，我使用了NVIDIA提供的一个非常还用的工具：torch2trt，它提供了类似Pytorch的API以加载并调用TensorRT模型，您可以使用以下命令来安装torch2trt：

pip install git+https://github.com/NVIDIA-AI-IOT/torch2trt.git

在真正读取TensorRT模型之前，你仍然需要安装TensorRT的Python API，它通常位于您TensorRT解压路径下的python目录下，请按照自己的Python版本按需安装。

接下来我们需要反序列化我们的YOLOv8 TensorRT模型：

import tensorrt as trt
from torch2trt import TRTModule
self.logger = trt.Logger(trt.Logger.INFO) # 初始化TensorRT日志类
# 反序列化YOLOv8模型
with open("../gpupipe/model/best.engine","rb") as f, trt.Runtime(self.logger) as runtime:
	self.engine = runtime.deserialize_cuda_engine(f.read())
# 使用torch2trt下TRTMoudle类来加载TensorRT引擎
self.TRTNet = TRTModule(input_names=[self.inputName],output_names=[self.outputName],engine=self.engine)

以上代码中：

input_name：代表您的YOLOv8模型的输入名称，如果您使用的是Ultralytics的工具进行训练的话，那么该值通常为：images。
output_name：代表您的YOLOv8模型的输出名称，如果您使用的是Ultralytics的工具进行训练的话，那么该值通常为：output0。
engine：代表反序列化后的TensorRT二进制对象。

正如我们前面有提到的一样，torch2trt为开发者提供了类Pytorch的API，因此我们的推理函数十分简单：

def inference(self,processedFrame:torch.Tensor) -> torch.Tensor:
    return self.TRTNet(processedFrame)

后处理函数

本章是整个pipeline最为重要的部分，这里我会着重讲解。
YOLOv8的输出大小通常取决于nc(num_classes)的大小，无论任何尺寸的YOLOv8模型最终都将输出8400条记录。这也就意味着我们需要遍历8400xnc个数据项！！这显然是没办法用循环来实现的。
在CPU环境上，我使用了Pandas进行向量化处理，得到了约9-10FPS美妙的成绩，具体实现代码如下：

output = np.transpose(np.squeeze(output[0])) # 由于YOLOv8的输出是(1,nc,8400)，为了方便处理所以需要转置
df = pd.DataFrame(output)
x_factor = self.imageWidth / self.inputWidth
y_factor = self.imageHeight / self.inputHeight

df['amax'] = np.amax(df.iloc[:,4:self.nc+4],axis=1)
df['argmax'] = np.where(df.iloc[:,4:self.nc+4]==df['amax'].values[:,None])[1]
df = df[df['amax'] > self.confidenceThres]
boxes = df.iloc[:,0:4].to_numpy()
scores = df['amax'].to_numpy()
class_ids = df['argmax'].to_numpy()

# apply factor to boxes
boxes[:,0] = (boxes[:,0] - boxes[:,2]/2.0) * x_factor
boxes[:,1] =