前言

RT-DETR(Real-Time Detetction Transformer) 是由 Baidu 提出的基于 transformer 的端到端实时检测器，本篇文章主要分享博主在实现 RT-DETR 推理和部署时做的一些尝试，不涉及任何的原理性分析。若有问题欢迎各位看官批评指正😄
参考：https://github.com/shouxieai/tensorRT_Pro
实现：https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8

1. RT-DETR-官方

博主先在 GitHub 上看到了官方 RT-DETR 的实现，它也提供了 pytorch 版本的实现并提供了导出 onnx 的工具，因此博主先拿它做的尝试。

RT-DETR-官方源码：https://github.com/lyuwenyu/RT-DETR/tree/main

官方在 rtdetr_pytorch 实现上也没有提供一个预测的代码，但博主在 rtdetr_pytorch/tools/export_onnx.py 文件中发现了一段利用 onnxruntime 来执行推理的代码，既然这样，那我们先来导出 onnx 模型，然后利用它的 onnxruntime 推理代码来执行推理

1. 源码下载

git clone https://github.com/lyuwenyu/RT-DETR.git

2. 权重准备，官方提供了如下的 pytorch 预训练权重：

博主选择了一个参数量最小的模型，即 rtdetr_r18vd 来完成后续的工作

将下载好的权重文件放在 RT-DETR/rtdetr_pytorch 文件夹下即可

也可以点击 here【pwd:yolo】 下载博主准备好的代码和权重文件（注意该代码和权重下载于 2023/11/11 日，若有改动请参考最新）

3. 导出 onnx 模型

按照官方的 README 文档来导出 onnx 模型，指令如下：

cd rtdetr_pytorch
python tools/export_onnx.py -c configs/rtdetr/rtdetr_r18vd_6x_coco.yml -r path/to/checkpoint --check --simplify

你可能会遇到如下问题：

提示说 ImportError: cannot import name ‘datapoints’ from ‘torchvision’

这是由于 torchvision 版本导致的问题，博主使用的是软件环境是 torch==2.1.0，torchvision==0.16.0；但是 requirements.txt 要求的环境是 torch==2.0.1，torchvision==0.15.2，因此博主换了一个虚拟环境，重新执行了上述导出代码。

在导出 onnx 模型时它还会去下载一个 ResNet18_vd_pretrained_from_paddle.pth 的预训练权重，导出过程如下：

可以看到导出时还是存在一些警告的，博主并未理会，执行成功后会在当前目录下生成 model.onnx 模型，我们可以使用 Netron 可视化工具看下导出的 onnx 模型，如下所示：

可以看到导出的 onnx 模型存在两个输入三个输出，我们再利用官方的 onnxruntime 推理代码看看能否推理成功，新建一个 predict.py 文件，其内容如下所示：

import torch
import onnxruntime as ort 
from PIL import Image, ImageDraw
from torchvision.transforms import ToTensor

if __name__ == "__main__":
    # print(onnx.helper.printable_graph(mm.graph))

    im = Image.open('bus.jpg').convert('RGB')
    im = im.resize((640, 640))
    im_data = ToTensor()(im)[None]
    print(im_data.shape)

    size = torch.tensor([[640, 640]])
    sess = ort.InferenceSession("model.onnx")
    output = sess.run(
        # output_names=['labels', 'boxes', 'scores'],
        output_names=None,
        input_feed={'images': im_data.data.numpy(), "orig_target_sizes": size.data.numpy()}
    )

    # print(type(output))
    # print([out.shape for out in output])

    labels, boxes, scores = output

    draw = ImageDraw.Draw(im)
    thrh = 0.6

    for i in range(im_data.shape[0]):

        scr = scores[i]
        lab = labels[i][scr > thrh]
        box = boxes[i][scr > thrh]

        print(i, sum(scr > thrh))

        for b in box:
            draw.rectangle(list(b), outline='red',)
            draw.text((b[0], b[1]), text=str(lab[i]), fill='blue', )

    im.save('test.jpg')

在终端执行如下指令即可完成推理：

python predict.py

执行成功后会在当前目录下保存推理的图片，如下所示：

可以看到推理还是没问题的，但是导出的 onnx 博主并不喜欢，首先它有两个输入，这就很头疼，其次它的三个输出也没有合并，还得去具体的代码中看看如何将它们合并，在后续部署时要填的坑比较多，因此博主果断放弃了这个方案。

虽然没有选择这个方案，但是博主在利用 onnxruntime 推理时也获取到了一些有用的信息，首先模型的预处理部分做了 BGR → RGB，ToTensor，添加 batch 维度等操作，但是似乎没有做 letterbox，而是直接进行的 resize；其次，后处理部分没有了 NMS，直接获取预测的每个目标的结果。

2. RT-DETR-U版

博主在 Ultralytics YOLOv8 的文档中看到 Ultralytics 也有 RT-DETR 的支持，因此打算再尝试下 U 版的 RT-DETR，虽然不如官方的原汁原味，但是作为学习还是可以的。

U 版提供了 RT-DETR 的推理、验证和训练，比较完善，还支持 onnx 导出，但它只提供 rtdetr-l.pt 和 rtdetr-x.pt 两个预训练权重，我们将使用 rtdetr-l.pt 完成我们后续的工作。

2.1 RT-DETR预测

我们先尝试利用官方预训练权重来推理一张图片并保存，看能否成功

在 YOLOv8 主目录下新建 predict.py 预测文件，其内容如下：

import cv2
from ultralytics import RTDETR

def hsv2bgr(h, s, v):
    h_i = int(h * 6)
    f = h * 6 - h_i
    p = v * (1 - s)
    q = v * (1 - f * s)
    t = v * (1 - (1 - f) * s)
    
    r, g, b = 0, 0, 0

    if h_i == 0:
        r, g, b = v, t, p
    elif h_i == 1:
        r, g, b = q, v, p
    elif h_i == 2:
        r, g, b = p, v, t
    elif h_i == 3:
        r, g, b = p, q, v
    elif h_i == 4:
        r, g, b = t, p, v
    elif h_i == 5:
        r, g, b = v, p, q

    return int(b * 255), int(g * 255), int(r * 255)

def random_color(id):
    h_plane = (((id << 2) ^ 0x937151) % 100) / 100.0
    s_plane = (((id << 3) ^ 0x315793) % 100) / 100.0
    return hsv2bgr(h_plane, s_plane, 1)

if __name__ == "__main__":

    model = RTDETR("rtdetr-l.pt")

    img = cv2.imread("ultralytics/assets/bus.jpg")
    results = model(img)[0]
    names   = model.names
    boxes   = results.boxes.data.tolist()

    for obj in boxes:
        left, top, right, bottom = int(obj[0]), int(obj[1]), int(obj[2]), int(obj[3])
        confidence = obj[4]
        label = int(obj[5])
        color = random_color(label)
        cv2.rectangle(img, (left, top), (right, bottom), color=color ,thickness=2, lineType=cv2.LINE_AA)
        caption = f"{names[label]} {confidence:.2f}"
        w, h = cv2.getTextSize(caption, 0, 1, 2)[0]
        cv2.rectangle(img, (left - 3, top - 33), (left + w + 10, top), color, -1)
        cv2.putText(img, caption, (left, top - 5), 0, 1, (0, 0, 0), 2, 16)

    cv2.imwrite("predict.jpg", img)
    print("save done")    

关于 RT-DETR 的预训练权重可以点击 here【pwd:yolo】 下载（注意该代码和权重下载于 2023/11/11 日，若有改动请参考最新）

在上述代码中我们通过 opencv 读取了一张图像，并送入模型中推理得到输出 results，results 中保存着不同任务的结果，我们这里是检测任务，因此只需要拿到对应的 boxes 即可。

拿到 boxes 后我们就可以将对应的框和模型预测的类别以及置信度绘制在图像上并保存。

关于可视化的代码实现参考自 tensorRT_Pro 中的实现，可以参考：app_yolo.cpp#L95

关于随机颜色的代码实现参考自 tensorRT_Pro 中的实现，可以参考：ilogger.cpp#L90

模型推理保存的结果图像如下所示：

2.2 RT-DETR预处理

模型预测成功后我们就需要自己动手来写下 RT-DETR 的预处理和后处理，方便后续在 C++ 上的实现，我们先来看看预处理的实现。

经过我们的调试分析可知 YOLOv8 的预处理过程在 ultralytics/engine/predictor.py 文件中，可以参考：predictor.py#L111

代码如下：

def preprocess(self, im):
    """
    Prepares input image before inference.

    Args:
        im (torch.Tensor | List(np.ndarray)): BCHW for tensor, [(HWC) x B] for list.
    """
    not_tensor = not isinstance(im, torch.Tensor)
    if not_tensor:
        im = np.stack(self.pre_transform(im))
        im = im[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW, (n, 3, h, w)
        im = np.ascontiguousarray(im)  # contiguous
        im = torch.from_numpy(im)

    im = im.to(self.device)
    im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
    if not_tensor:
        im /= 255  # 0 - 255 to 0.0 - 1.0
    return im

它包含以下步骤：

• self.pre_transform：注意这里最终会调用 rtdetr/predict.py#L71，和 YOLO 检测器的 letterbox 操作不同，它其实就是做了一个 resize 的操作
• im[…,::-1]：BGR → RGB
• transpose((0, 3, 1, 2))：添加 batch 维度，HWC → CHW
• torch.from_numpy：to Tensor
• im /= 255：除以 255，归一化

因此我们不难写出对应的预处理代码，如下所示：

def preprocess(image):
    image = cv2.resize(image, (640, 640))
    image = (image[..., ::-1] / 255.0).astype(np.float32) # BGR to RGB, 0 - 255 to 0.0 - 1.0
    image = image.transpose(2, 0, 1)[None]  # BHWC to BCHW (n, 3, h, w)
    image = torch.from_numpy(image)
    return image

2.3 RT-DETR后处理

其实 RT-DETR 没有后处理一说，因为它是端到端的检测器，因此只需要我们根据置信度阈值过滤框即可

经过我们的调试分析可知 RT-DETR 的后处理过程在 ultralytics/models/rtdetr/predict.py 文件中，可以参考：rtdetr/predict.py#L34

class RTDETRPredictor(BasePredictor):
    """
    RT-DETR (Real-Time Detection Transformer) Predictor extending the BasePredictor class for making predictions using
    Baidu's RT-DETR model.

    This class leverages the power of Vision Transformers to provide real-time object detection while maintaining
    high accuracy. It supports key features like efficient hybrid encoding and IoU-aware query selection.

    Example:
        ```python
        from ultralytics.utils import ASSETS
        from ultralytics.models.rtdetr import RTDETRPredictor

        args = dict(model='rtdetr-l.pt', source=ASSETS)
        predictor = RTDETRPredictor(overrides=args)
        predictor.predict_cli()
        ```

    Attributes:
        imgsz (int): Image size for inference (must be square and scale-filled).
        args (dict): Argument overrides for the predictor.
    """

    def postprocess(self, preds, img, orig_imgs):
        """
        Postprocess the raw predictions from the model to generate bounding boxes and confidence scores.

        The method filters detections based on confidence and class if specified in `self.args`.

        Args:
            preds (torch.Tensor): Raw predictions from the model.
            img (torch.Tensor): Processed input images.
            orig_imgs (list or torch.Tensor): Original, unprocessed images.

        Returns:
            (list[Results]): A list of Results objects containing the post-processed bounding boxes, confidence scores,
                and class labels.
        """
        nd = preds[0].shape[-1]
        bboxes, scores = preds[0].split((4, nd - 4), dim=-1)

        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

        results = []
        for i, bbox in enumerate(bboxes):  # (300, 4)
            bbox = ops.xywh2xyxy(bbox)  # (300, 4)
            score, cls = scores[i].max(-1, keepdim=True)  # (300, 1)
            idx = score.squeeze(-1) > self.args.conf  # (300, )
            if self.args.classes is not None:
                idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
            pred = torch.cat([bbox, score, cls], dim=-1)[idx]  # filter
            orig_img = orig_imgs[i]
            oh, ow = orig_img.shape[:2]
            pred[..., [0, 2]] *= ow
            pred[..., [1, 3]] *= oh
            img_path = self.batch[0][i]
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
        return results

它包含以下步骤：

• xywh2xyxy：框的转换
• score.squeeze(-1) > self.args.conf：框的过滤
• pred[…, [0, 2]] *= ow：框坐标映射

后处理部分非常简单，因此我们不难写出对应的后处理代码，如下所示：

def postprocess(pred, oh, ow, conf_thres=0.25):

    # 输入是模型推理的结果，即300个预测框
    # 1,300,84 [cx,cy,w,h,class*80]
    boxes = []
    for item in pred[0]:
        cx, cy, w, h = item[:4]
        label = item[4:].argmax()
        confidence = item[4 + label]
        if confidence < conf_thres:
            continue
        left    = cx - w * 0.5
        top     = cy - h * 0.5
        right   = cx + w * 0.5
        bottom  = cy + h * 0.5
        boxes.append([left, top, right, bottom, confidence, label])

    boxes = np.array(boxes)
    lr = boxes[:,[0, 2]]
    tb = boxes[:,[1, 3]]
    boxes[:,[0,2]] = ow * lr
    boxes[:,[1,3]] = oh * tb

    return boxes

对于一张 640x640 的图片来说，RT-DETR 预测框的总数量是 300，每个预测框的维度是 84（针对 COCO 数据集的 80 个类别而言）
$$300\times 84=300\times (4+80)$$
其中的 4 对应的是 cx，cy，w，h，分别代表的含义是边界框中心点坐标、宽高；80 对应的是 COCO 数据集中的 80 个类别置信度。

2.4 RT-DETR推理

通过上面对 RT-DETR 的预处理和后处理分析之后，整个推理过程就显而易见了。RT-DETR 的推理包括图像预处理、模型推理、预测结果后处理三部分，其中预处理主要包括 resize，后处理主要包括框的变换。

完整的推理代码如下：

import cv2
import torch
import numpy as np
from ultralytics.nn.autobackend import AutoBackend

def preprocess(image):
    image = cv2.resize(image, (640, 640))
    image = (image[..., ::-1] / 255.0).astype(np.float32) # BGR to RGB, 0 - 255 to 0.0 - 1.0
    image = image.transpose(2, 0, 1)[None]  # BHWC to BCHW (n, 3, h, w)
    image = torch.from_numpy(image)
    return image

def postprocess(pred, oh, ow, conf_thres=0.25):

    # 输入是模型推理的结果，即300个预测框
    # 1,300,84 [cx,cy,w,h,class*80]
    boxes = []
    for item in pred[0]:
        cx, cy, w, h = item[:4]
        label = item[4:].argmax()
        confidence = item[4 + label]
        if confidence < conf_thres:
            continue
        left    = cx - w * 0.5
        top     = cy - h * 0.5
        right   = cx + w * 0.5
        bottom  = cy + h * 0.5
        boxes.append([left, top, right, bottom, confidence, label])

    boxes = np.array(boxes)
    lr = boxes[:,[0, 2]]
    tb = boxes[:,[1, 3]]
    boxes[:,[0,2]] = ow * lr
    boxes[:,[1,3]] = oh * tb

    return boxes

def hsv2bgr(h, s, v):
    h_i = int(h * 6)
    f = h * 6 - h_i
    p = v * (1 - s)
    q = v * (1 - f * s)
    t = v * (1 - (1 - f) * s)
    
    r, g, b = 0, 0, 0

    if h_i == 0:
        r, g, b = v, t, p
    elif h_i == 1:
        r, g, b = q, v, p
    elif h_i == 2:
        r, g, b = p, v, t
    elif h_i == 3:
        r, g, b = p, q, v
    elif h_i == 4:
        r, g, b = t, p, v
    elif h_i == 5:
        r, g, b = v, p, q

    return int(b * 255), int(g * 255), int(r * 255)

def random_color(id):
    h_plane = (((id << 2) ^ 0x937151) % 100) / 100.0
    s_plane = (((id << 3) ^ 0x315793) % 100) / 100.0
    return hsv2bgr(h_plane, s_plane, 1)

if __name__ == "__main__":
    
    img = cv2.imread("ultralytics/assets/bus.jpg")
    oh, ow = img.shape[:2]

    img_pre = preprocess(img)

    # postprocess
    # ultralytics/models/rtdetr/predict.py
    model  = AutoBackend(weights="rtdetr-l.pt")
    names  = model.names
    result = model(img_pre)[0]  # 1,300,84

    boxes  = postprocess(result, oh, ow)

    for obj in boxes:
        left, top, right, bottom = int(obj[0]), int(obj[1]), int(obj[2]), int(obj[3])
        confidence = obj[4]
        label = int(obj[5])
        color = random_color(label)
        cv2.rectangle(img, (left, top), (right, bottom), color=color ,thickness=2, lineType=cv2.LINE_AA)
        caption = f"{names[label]} {confidence:.2f}"
        w, h = cv2.getTextSize(caption, 0, 1, 2)[0]
        cv2.rectangle(img, (left - 3, top - 33), (left + w + 10, top), color, -1)
        cv2.putText(img, caption, (left, top - 5), 0, 1, (0, 0, 0), 2, 16)

    cv2.imwrite("infer.jpg", img)
    print("save done")  

推理效果如下图所示：

至此，我们在 Python 上面完成了 RT-DETR 的整个推理过程，下面我们去 C++ 上实现。

3. RT-DETR-C++

C++ 上的实现我们使用的是 repo 依旧是 tensorRT_Pro，现在我们就基于 tensorRT_Pro 完成 RT-DETR 在 C++上的推理。

3.1 ONNX导出

首先我们需要将 RT-DETR 模型导出为 ONNX，为了适配 tensorRT_Pro 我们需要做一些修改，主要有以下几点：

• 修改输出节点名为 output
• 输入输出只让 batch 维度动态，宽高不动态

具体修改如下：

1. 在 ultralytics/engine/exporter.py 文件中改动一处

• 323 行：输出节点名修改为 output
• 326 行：输入只让 batch 维度动态，宽高不动态
• 331 行：输出只让 batch 维度动态，宽高不动态

# ========== exporter.py ==========

# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
#     dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
#     if isinstance(self.model, SegmentationModel):
#         dynamic['output0'] = {0: 'batch', 2: 'anchors'}  # shape(1, 116, 8400)
#         dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
#     elif isinstance(self.model, DetectionModel):
#         dynamic['output0'] = {0: 'batch', 2: 'anchors'}  # shape(1, 84, 8400)
# 修改为：

output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:
    dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)
    if isinstance(self.model, SegmentationModel):
        dynamic['output0'] = {0: 'batch', 2: 'anchors'}  # shape(1, 116, 8400)
        dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
    elif isinstance(self.model, DetectionModel):
        dynamic['output'] = {0: 'batch'}  # shape(1, 84, 8400)

以上就是为了适配 tensorRT_Pro 而做出的代码修改，修改好以后，将预训练权重 rtdetr-l.pt 放在 ultralytics-main 主目录下，新建导出文件 export.py，内容如下：

from ultralytics import RTDETR

model = RTDETR("rtdetr-l.pt")

success = model.export(format="onnx", dynamic=True, simplify=True)

在终端执行如下指令即可完成 onnx 导出：

python export.py

你可能会遇到如下的问题：

提示说 Unsupported: ONNX export of operator get_pool_ceil_padding

在 ultralytics/issues/6144 你会发现有人存在相同的问题，作者最终将问题定位在 torch 版本问题，因此博主尝试替换了一个虚拟环境，博主最初的 torch 版本是 2.1.0，将其替换成 2.0.1 之后就没有问题了。

导出过程如下图所示：

可以看到导出的 pytorch 模型的输入 shape 是 1x3x640x640，输出 shape 是 1x300x84，符合我们的预期。

导出成功后会在当前目录下生成 rtdetr-l.onnx 模型，我们可以使用 Netron 可视化工具查看，如下图所示：

可以看到输入节点名是 images，维度是 batchx3x640x640，保证只有 batch 维度动态，但是让博主困惑的是输出节点名是 output，维度却是 1x300x84，没有看到动态 batch 维度，但是在后续测试部署过程中并没有发现什么问题。

大家如果担心出什么问题的话，可以将 dynamic 参数设置为 False，导出静态 onnx 模型也行。

3.2 RT-DETR预处理

之前有提到过 RT-DETR 的预处理部分就是 resize，而在 tensorRT_Pro 中有提供 CUDA 版本的 resize 实现，我们直接拿过来使用即可。

tensorRT_Pro 中预处理的代码如下：

__global__ void resize_bilinear_and_normalize_kernel(
    uint8_t* src, int src_line_size, int src_width, int src_height, float* dst, int dst_width, int dst_height, 
    float sx, float sy, Norm norm, int edge
){
    int position = blockDim.x * blockIdx.x + threadIdx.x;
    if (position >= edge) return;

    int dx      = position % dst_width;
    int dy      = position / dst_width;
    float src_x = (dx + 0.5f) * sx - 0.5f;
    float src_y = (dy + 0.5f) * sy - 0.5f;
    float c0, c1, c2;

    int y_low = floorf(src_y);
    int x_low = floorf(src_x);
    int y_high = limit(y_low + 1, 0, src_height - 1);
    int x_high = limit(x_low + 1, 0, src_width - 1);
    y_low = limit(y_low, 0, src_height - 1);
    x_low = limit(x_low, 0, src_width - 1);

    int ly    = rint((src_y - y_low) * INTER_RESIZE_COEF_SCALE);
    int lx    = rint((src_x - x_low) * INTER_RESIZE_COEF_SCALE);
    int hy    = INTER_RESIZE_COEF_SCALE - ly;
    int hx    = INTER_RESIZE_COEF_SCALE - lx;
    int w1    = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
    float* pdst = dst + dy * dst_width + dx * 3;
    uint8_t* v1 = src + y_low * src_line_size + x_low * 3;
    uint8_t* v2 = src + y_low * src_line_size + x_high * 3;
    uint8_t* v3 = src + y_high * src_line_size + x_low * 3;
    uint8_t* v4 = src + y_high * src_line_size + x_high * 3;

    c0 = resize_cast(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]);
    c1 = resize_cast(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]);
    c2 = resize_cast(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]);

    if(norm.channel_type == ChannelType::Invert){
        float t = c2;
        c2 = c0;  c0 = t;
    }

    if(norm.type == NormType::MeanStd){
        c0 = (c0 * norm.alpha - norm.mean[0]) / norm.std[0];
        c1 = (c1 * norm.alpha - norm.mean[1]) / norm.std[1];
        c2 = (c2 * norm.alpha - norm.mean[2]) / norm.std[2];
    }else if(norm.type == NormType::AlphaBeta){
        c0 = c0 * norm.alpha + norm.beta;
        c1 = c1 * norm.alpha + norm.beta;
        c2 = c2 * norm.alpha + norm.beta;
    }

    int area = dst_width * dst_height;
    float* pdst_c0 = dst + dy * dst_width + dx;
    float* pdst_c1 = pdst_c0 + area;
    float* pdst_c2 = pdst_c1 + area;
    *pdst_c0 = c0;
    *pdst_c1 = c1;
    *pdst_c2 = c2;
}

关于预处理部分其实就是调用了上述 CUDA 核函数来实现 warpAffine，由于在 CUDA 中我们是对每个像素进行操作，因此非常容易实现 BGR → RGB，/255.0 等操作。

3.3 RT-DETR后处理

之前有提到 RT-DETR 是基于端到端的检测器，是没有后处理的，不过我们还是需要将框进行 decode 解码，代码可参考：yolo_decode.cu#L13

因此我们不难写出 RT-DETR 的 decode 解码部分的实现代码，如下所示：

static __global__ void decode_kernel(float *predict, int num_bboxes, int num_classes, float confidence_threshold, float* parray, int MAX_IMAGE_BOXES){
    
    int position = blockDim.x * blockIdx.x + threadIdx.x;
    if (position >= num_bboxes) return;

    float* pitem            = predict + (4 + num_classes) * position;
    float* class_confidence = pitem + 4;
    float confidence        = *class_confidence++;
    int label               = 0;
    for(int i = 1; i < num_classes; ++i, ++class_confidence){
        if(*class_confidence > confidence){
            confidence = *class_confidence;
            label      = i;
        }
    }

    if(confidence < confidence_threshold)
        return;

    int index = atomicAdd(parray, 1);
    if(index >= MAX_IMAGE_BOXES)
        return;

    float cx         = *pitem++;
    float cy         = *pitem++;
    float width      = *pitem++;
    float height     = *pitem++;
    float left   = cx - width  * 0.5f;
    float top    = cy - height * 0.5f;
    float right  = cx + width  * 0.5f;
    float bottom = cy + height * 0.5f;

    float *pout_item = parray + 1 + index * NUM_BOX_ELEMENT;
    *pout_item++ = left;
    *pout_item++ = top;
    *pout_item++ = right;
    *pout_item++ = bottom;
    *pout_item++ = confidence;
    *pout_item++ = label;
}

关于 decode 的具体实现其实就是启动多个线程，每个线程处理一个框的解码，由于解码出来的框坐标是在 640x640 的图像上，因此可视化时还需要将其映射到原图上。

3.4 RT-DETR推理

通过上面对 RT-DETR 的预处理和后处理分析之后，整个推理过程就显而易见了。C++ 上 RT-DETR 的预处理部分沿用 CUDA 版本的 resize，后处理中的 decode 解码部分简单修改即可。

首先我们需要利用 tensorRT_Pro 的编译接口 TRT::compile 来将 RT-DETR 的 ONNX 模型生成对应版本的 engine，编译图解如下所示：

可以看到提示 Pad 节点的解析存在问题，我们可以去 onnx_parser/builtin_op_importers.cpp 文件中搜索下看是否支持 Pad 节点，在 builtin_op_importers.cpp#L3028 中可以看到 onnx_parser 解析器是支持 Pad 节点的解析的，但是依旧解析错误，说明 RT-DETR 的 Pad 的实现和 onnx_parser 解析的 Pad 存在一定的出入，毕竟 tensorRT_Pro 中的 onnx-parser 解析器是 8.0 版本的，有点老了。

我们目前无法通过 TRT::compile 编译接口生成 engine，博主想到了两种方案，一种是在 tensorRT高性能部署课程 中杜老师教过的替换 onnx-parser 解析器，替换成高版本的 onnx-parser 解析器再看是否存在节点解析问题；另外一种就是利用高版本 tensorRT 的 trtexec 工具生成 engine

博主先采用的第二种方案，利用高版本的 trtexec 工具生成 engine

博主新建了一个 build.sh 脚本文件，其内容如下：

#! /usr/bin/bash

TRTEXEC=/opt/TensorRT-8.4.1.5/bin/trtexec

${TRTEXEC} --onnx=rtdetr-l.onnx --minShapes=images:1x3x640x640 --optShapes=images:1x3x640x640 --maxShapes=images:16x3x640x640 --saveEngine=rtdetr-l.FP32.trtmodel

在终端执行如下指令即可：

bash build.sh

输出如下：

可以看到 trtexec 工具生成 engine 也失败了，提示说 LayerNormalization 不支持，至少 Pad 节点解析的问题解决了，只不过又出了新的节点解析问题，博主目前使用的 tensorRT 的版本是 8.4.1.5，难道需要自己写插件支持嘛

博主到 onnx-tensorrt 官网的主分支下搜索了下想看看最新的 onnx-parser 解析器是否支持 LayerNormalization 层，发现竟然支持，既然这样没必要自己实现插件了呀，具体可参考：builtin_op_importers.cpp#L2270

只是博主安装的 tensorRT 版本太低了，还不支持 LayerNormalization 节点的解析，经博主研究发现只有在最新 release/8.6-GA 版本才有 LayerNormalization 层的支持，因此博主又安装了一个最新版本的 tensorRT

关于 tensorRT 的安装可以参考：Ubuntu20.04软件安装大全

记得配置下环境变量，不然仍可能报错

外网访问较慢，这边也提供博主下载好的安装包，点击 here【pwd:yolo】 下载即可

安装完成后再修改重新指定下 trtexec 的路径为最新的 tensorRT-8.6.1.6 的路径即可，再次执行 build.sh 文件输出如下：

可以看到 engine 生成成功了，接下来就是拿着 engine 去进行推理了

我们在终端执行如下指令即可完成推理（注意！完整流程博主会在后续内容介绍，这边只是简单演示）

make rtdetr -j64

编译图解如下所示：

可以看到叒报错了，报错信息提示我们在反序列化模型的时候出现了问题，这是什么原因导致的呢？🤔

噢！博主想起来了，博主在 Makefile 中指定链接的 tensorRT 的库文件还是 8.4.1 版本的，而生成 engine 的tensorRT 版本是 8.6.1 版本的，序列化 engine 和反序列化 engine 的 tensorRT 不是同一个版本，肯定会报错呀！因此你需要在 Makefile 中重新修改下 tensorRT 的路径指定，如下所示：

# RT-DETR 必须指定高版本的 tensorRT
lean_tensor_rt := /home/jarvis/lean/TensorRT-8.6.1.6

先执行 make clean 清除下编译文件，然后再去执行 make rtdetr -j64 可以看到输出如下：

总算是推理成功了，不容易吖😂

推理结果如下图所示：

至此，我们在 C++ 上面完成了 RT-DETR 的整个推理过程，下面我们将完整的走一遍流程。

4. RT-DETR部署

博主新建了一个仓库 tensorRT_Pro-YOLOv8，该仓库基于 shouxieai/tensorRT_Pro，并进行了调整以支持 YOLOv8 的各项任务，目前已支持分类、检测、分割、姿态点估计任务。

下面我们就来具体看看如何利用 tensorRT_Pro-YOLOv8 这个 repo 完成 RT-DETR 的推理。

4.1 源码下载

tensorRT_Pro-YOLOv8 的代码可以直接从 GitHub 官网上下载，源码下载地址是 https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8，Linux 下代码克隆指令如下：

git clone https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2023/11/11 日，若有改动请参考最新）

4.2 环境配置

需要使用的软件环境有 TensorRT、CUDA、cuDNN、OpenCV、Protobuf，所有软件环境的安装可以参考 Ubuntu20.04软件安装大全，这里不再赘述，需要各位看官自行配置好相关环境😄，外网访问较慢，这里提供下博主安装过程中的软件安装包下载链接 Baidu Drive【pwd:yolo】🚀🚀🚀

tensorRT_Pro-YOLOv8 提供 CMakeLists.txt 和 Makefile 两种方式编译，二者选一即可

4.2.1 配置CMakeLists.txt

主要修改五处

1. 修改第 13 行，修改 OpenCV 路径

set(OpenCV_DIR   "/usr/local/include/opencv4/")

2. 修改第 15 行，修改 CUDA 路径

set(CUDA_TOOLKIT_ROOT_DIR     "/usr/local/cuda-11.6")

3. 修改第 16 行，修改 cuDNN 路径

set(CUDNN_DIR    "/usr/local/cudnn8.4.0.27-cuda11.6")

4. 修改第 17 行，修改 tensorRT 路径（版本必须大于 8.6）

set(TENSORRT_DIR "/home/jarvis/lean/TensorRT-8.6.1.6")

5. 修改第 20 行，修改 protobuf 路径

set(PROTOBUF_DIR "/home/jarvis/protobuf")

4.2.2 配置Makefile

主要修改五处

1. 修改第 4 行，修改 protobuf 路径

lean_protobuf  := /home/jarvis/protobuf

2. 修改第 5 行，修改 tensorRT 路径（版本必须大于 8.6）

lean_tensor_rt := /home/jarvis/lean/TensorRT-8.6.1.6

3. 修改第 6 行，修改 cuDNN 路径

lean_cudnn     := /usr/local/cudnn8.4.0.27-cuda11.6

4. 修改第 7 行，修改 OpenCV 路径

lean_opencv    := /usr/local

5. 修改第 8 行，修改 CUDA 路径

lean_cuda      := /usr/local/cuda-11.6

4.3 ONNX导出

导出细节可以查看之前的内容，这边不再赘述。记得将导出的 ONNX 模型放在 tensorRT_Pro-YOLOv8/workspace 文件夹下。

4.4 engine生成

修改 workspace 下 build.sh 文件内容，修改 trtexec 的路径为你自己的路径，终端执行如下指令：

cd tensorRT_Pro-YOLOv8/workspace
bash build.sh

4.5 源码修改

如果你想推理自己训练的模型还需要修改下源代码，RT-DETR 模型的推理代码主要在 app_rtdetr.cpp 文件中，我们就只需要修改这一个文件中的内容即可，源码修改较简单主要有以下几点：

• 1. app_rtdetr.cpp 268行，“rtdetr-l” 修改为你导出的 ONNX 模型名
• 2. app_rtdetr.cpp 10行，将 cocolabels 数组中的类别名称修改为你训练的类别

具体修改示例如下：

test(TRT::Mode::FP32, "best")	// 修改1 268行"rtdetr-l"改成"best"

static const char *cocolabels[] = {"have_mask", "no_mask"};	// 修改2 10行修改检测类别，为自训练模型的类别名称

4.6 运行

OK！源码修改好了，Makefile 编译文件也搞定了，ONNX 模型也准备好了，现在可以编译运行了，直接在终端执行如下指令即可：

make rtdetr -j64

编译过程如下所示：

编译运行成功后会生成 rtdetr-l_RT-DETR_FP32_result 文件夹，该文件夹下保存了推理的图片。

模型推理效果如下图所示：

OK！以上就是使用 tensorRT_Pro-YOLOv8 推理 RT-DETR 的大致流程，若有问题，欢迎各位看官批评指正。

5. 拓展-onnx-tensorrt配置

之前我们不是说 RT-DETR 的 engine 生成有两种方案嘛，一种是自己下载编译配置 onnx-tensorrt，自己构建 onnx-parser 解析器，另一种是直接通过 tensorRT 官方的 libnvonnxparser.so 来解析 ONNX 文件，上面我们是通过第二种方法来实现 engine 的生成的，这里我们通过第一种方法来生成 engine，再来回顾下杜老师之前教过的知识

在开始之前你依旧需要安装 tensorRT-8.6.1 这个高版本的 tensorRT，这是因为 tensorRT 是跟 onnx 解析器挂钩的，我们在这里相当于只是自己来构建 libnvonnxpaser.so，不使用官方提供的版本，因此要版本匹配。

我们先看下 tensorRT 是什么版本，我们进入到 tensorRT 的安装目录，它的 include 文件夹下有一个 NvInferVersion.h 文件，将其打印出来，如下图所示：

从上图可知博主的 tensorRT 版本为 8.6.1，所以我们选择 onnx-tensorrt-8.6-GA 版本，其实我们通过之前的分析也知道只有这个版本的解析器才支持 LayerNormalization 算子的解析

onnx-tensorrt 有一个 third_party 的第三方库，打开其实就是 onnx，本质就是一个套娃，我们只要知道 onnx-tensorrt 依赖自 onnx 就行了

我们在其 README 文档中可以看到其安装要求，如下所示

我们先把它下载下来后再去使用它

下载地址：https://github.com/onnx/onnx-tensorrt/tree/release/8.6-GA

也可以点击 here【pwd:yolo】 下载博主准备好的源代码

由于实际替换过程有些繁琐,因此博主在这里就不一一说明具体修改的原因了,只讲解如何实现替换.博主也是对照着之前杜老师的视频,走一步看一步,错一步改一步来完成的

Step 1. 解压,删除不必要的文件

删除后的剩余文件如下所示:

Step 2. ImporterContext.hpp 注释第 10 行

// #include "onnx/common/stl_backports.h"

Step 3. ImporterContext.hpp 修改第 121 行

ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
//     : mNetwork(network)
//     , mLogger(logger)
//     , mErrorWrapper(onnx::make_unique<ErrorRecorderWrapper>(mNetwork, logger))
// {
// }
    
// 修改为:
    
ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
    : mNetwork(network)
    , mLogger(logger)
    , mErrorWrapper(nullptr)
{
}

Step 4. build_op_importers.cpp 新增头文件, 28 行新增函数

#include <onnxplugin/onnxplugin.hpp>


// 28 行新增函数
typedef std::function<std::vector<int64_t>(const std::string& name, const std::vector<int64_t>& shape)> layerhook_func_reshape;

static layerhook_func_reshape g_layerhook_func_reshape;
extern "C" TENSORRTAPI void register_layerhook_reshape(const layerhook_func_reshape& func){
    g_layerhook_func_reshape = func;
}

// 173 行新增函数
namespace onnx2trt
{
    ...
        
static TRT::DataType convert_trt_datatype(::onnx::TensorProto::DataType dt){
    switch(dt){
        case ::onnx::TensorProto::FLOAT: return TRT::DataType::Float;
        case ::onnx::TensorProto::FLOAT16: return TRT::DataType::Float16;
        case ::onnx::TensorProto::INT32: return TRT::DataType::Int32;
        case ::onnx::TensorProto::UINT8: return TRT::DataType::UInt8;
        default:
            printf("Unsupport data type %d\n", dt);
            return TRT::DataType::Unknow;
    }
}

DEFINE_BUILTIN_OP_IMPORTER(Plugin)
{
    std::vector<nvinfer1::ITensor*> inputTensors;
    std::vector<onnx2trt::ShapedWeights> weights;
    for(int i = 0; i < inputs.size(); ++i){
        auto& item = inputs.at(i);
        if(item.is_tensor()){
            nvinfer1::ITensor* input = &convertToTensor(item, ctx);
            inputTensors.push_back(input);
        }else{
            weights.push_back(item.weights());
        }
    }

    OnnxAttrs attrs(node, ctx);
    auto name = attrs.get<std::string>("name", "");
    auto info = attrs.get<std::string>("info", "");

    // Create plugin from registry
    auto registry = getPluginRegistry();
    auto creator = registry->getPluginCreator(name.c_str(), "1", "");
    if(creator == nullptr){
        printf("%s plugin was not found in the plugin registry!", name.c_str());
        ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);
    }
    
    nvinfer1::PluginFieldCollection pluginFieldCollection;
    pluginFieldCollection.nbFields = 0;

    ONNXPlugin::TRTPlugin* plugin = (ONNXPlugin::TRTPlugin*)creator->createPlugin(name.c_str(), &pluginFieldCollection);
    if(plugin == nullptr){
        LOG_ERROR(name << " plugin was not found in the plugin registry!");
        ASSERT(false, ErrorCode::kUNSUPPORTED_NODE);
    }

    std::vector<std::shared_ptr<TRT::Tensor>> weightTensors;
    for(int i = 0; i < weights.size(); ++i){
        auto& weight = weights[i];
        std::vector<int> dims(weight.shape.d, weight.shape.d + weight.shape.nbDims);
        auto onnx_dtype = convert_trt_datatype((::onnx::TensorProto::DataType)weight.type);
        if(onnx_dtype == TRT::DataType::Unknow){
            LOG_ERROR("unsupport weight type: " << weight.type);
        }
        
        std::shared_ptr<TRT::Tensor> dweight(new TRT::Tensor(dims, onnx_dtype));
        memcpy(dweight->cpu(), weight.values, dweight->bytes());
        weightTensors.push_back(dweight);
    }
    
    plugin->pluginInit(name, info, weightTensors);
    auto layer = ctx->network()->addPluginV2(inputTensors.data(), inputTensors.size(), *plugin);
    std::vector<TensorOrWeights> outputs;
    for( int i=0; i< layer->getNbOutputs(); ++i )
      outputs.push_back(layer->getOutput(i));
    return outputs;
}    
    
    ...
}

Step 5. 命名空间替换,将所有文件下的 ONNX_NAMESPACE 命名空间替换为 onnx

Step 6. NvOnnxParser.h 新增头文件,326 行新增函数

#include <functional>
#include <string>


// 326 行新增函数
extern "C" TENSORRTAPI void register_layerhook_reshape(const std::function<std::vector<int64_t>(const std::string& name, const std::vector<int64_t>& shape)>&);

Step 7. src/tensorRT/builder/trt_builder.cpp 修改 511 行

// 511 行修改
// onnxParser.reset(nvonnxparser::createParser(*network, gLogger, dims_setup), destroy_nvidia_pointer<nvonnxparser::IParser>);
// 修改为:

onnxParser.reset(nvonnxparser::createParser(*network, gLogger), destroy_nvidia_pointer<nvonnxparser::IParser>);

Step 8. ModelImporter.cpp 784 行新增函数

// 784 行新增
bool ModelImporter::parseFromData(const void* onnx_data, size_t size, int verbosity)
{
    GOOGLE_PROTOBUF_VERIFY_VERSION;
    ::onnx::ModelProto onnx_model;
    auto* ctx = &mImporterCtx;

    if (onnx_data == nullptr || size < 1)
    {
        LOG_ERROR("Failed to parse ONNX model from data, ptr = " << onnx_data << ", size = " << size);
        return false;
    }

    // Keep track of the absolute path to the ONNX file.
    const int64_t opset_version = (onnx_model.opset_import().size() ? onnx_model.opset_import(0).version() : 0);
    LOG_INFO("----------------------------------------------------------------");
    LOG_INFO("Input data size:   " << size);
    LOG_INFO("ONNX IR version:  " << onnx_ir_version_string(onnx_model.ir_version()));
    LOG_INFO("Opset version:    " << opset_version);
    LOG_INFO("Producer name:    " << onnx_model.producer_name());
    LOG_INFO("Producer version: " << onnx_model.producer_version());
    LOG_INFO("Domain:           " << onnx_model.domain());
    LOG_INFO("Model version:    " << onnx_model.model_version());
    LOG_INFO("Doc string:       " << onnx_model.doc_string());
    LOG_INFO("----------------------------------------------------------------");

    { //...Read input file, parse it
        if (!parse(onnx_data, size))
        {
            const int32_t nerror = getNbErrors();
            for (int32_t i = 0; i < nerror; ++i)
            {
                nvonnxparser::IParserError const* error = getError(i);
                if (error->node() != -1)
                {
                    ::onnx::NodeProto const& node = onnx_model.graph().node(error->node());
                    LOG_ERROR("While parsing node number " << error->node() << " [" << node.op_type() << " -> \"" << node.output(0) << "\"" << "]:");
                    LOG_ERROR("--- Begin node ---");
                    LOG_ERROR(pretty_print_onnx_to_string(node));
                    LOG_ERROR("--- End node ---");
                }
                LOG_ERROR("ERROR: " << error->file() << ":" << error->line() << " In function " << error->func() << ":\n"
                     << "[" << static_cast<int>(error->code()) << "] " << error->desc());
            }
            return false;
        }
    } //...End Reading input file, parsing it
    return true;
}

Step 9. ModelImporter.hpp 92 行新增

// 92 行新增
bool parseFromData(const void* onnx_data, size_t size, int verbosity) override;

Step 10. NvOnnxParser.h 184 行新增

// 184 行新增
virtual bool parseFromData(const void* onnx_data, size_t size, int verbosity) = 0;

Step 11. Makefile 将 C++ 标准修改为 C++14

# cpp_compile_flags := -std=c++11 -g -w -O0 -fPIC -pthread -fopenmp
# cu_compile_flags  := -std=c++11 -g -w -O0 -Xcompiler "$(cpp_compile_flags)" $(cuda_arch)
# 修改为:

cpp_compile_flags := -std=c++14 -g -w -O0 -fPIC -pthread -fopenmp
cu_compile_flags  := -std=c++14 -g -w -O0 -Xcompiler "$(cpp_compile_flags)" $(cuda_arch)

OK！以上就是全部的修改内容了

完整的文件内容可以参考：onnx_parser/onnx_parser_8.6

修改完成后我们新建一个 use_tensorrt_8.6.sh 脚本文件，其内容如下：

#!/bin/bash

echo Remove src/tensorRT/onnx_parser
rm -rf src/tensorRT/onnx_parser

echo Copy [onnx_parser/onnx-tensorrt-release-8.6-GA] to [src/tensorRT/onnx_parser]
cp -r onnx_parser/onnx-tensorrt-release-8.6-GA src/tensorRT/onnx_parser

echo Configure your tensorRT path to 8.6
echo After that, you can execute the command 'make rtdetr -j64'

在终端执行如下指令即可完成 onnx-parser 的替换

bash onnx_parser/use_tensorrt_8.6.sh

替换完成后我们就可以愉快的使用 TRT::Compile 接口来编译模型了，编译过程如下图所示：

编译运行成功后在 workspace 文件夹下会生成 engine 文件 rtdetr-l.FP32.trtmodel 用于模型推理，同时它还会生成 rtdetr-l_RT-DETR_FP32_result 文件夹，该文件夹下保存了推理的图片。

模型推理结果如下图所示：

OK！以上就是配置 onnx-tensorrt 手动替换 onnx-parser 解析器的大致流程，若有问题，欢迎各位看官批评指正。

结语

博主在这里针对 RT-DETR 的预处理和后处理做了简单分析，同时与大家分享了 C++ 上的实现流程，目的是帮大家理清思路，更好的完成后续的部署工作😄。感谢各位看到最后，创作不易，读后有收获的看官请帮忙点个👍⭐️

最后大家如果觉得 tensorRT_Pro-YOLOv8 这个 repo 对你有帮助的话，不妨点个 ⭐️ 支持一波，这对博主来说非常重要，感谢各位🙏。

下载链接

• 软件安装包下载链接【提取码:yolo】🚀🚀🚀
• 源代码、权重下载链接【提取码:yolo】

参考

• https://github.com/lyuwenyu/RT-DETR
• https://github.com/onnx/onnx-tensorrt
• https://github.com/ultralytics/ultralytics
• https://github.com/shouxieai/tensorRT_Pro