SSD目标检测网络tensorRT推理【附代码】

终于更新了，本篇是实现了SSD的tensorrt 推理【python版】。YOLOv4以及YOLOv5C++版的tensorrt推理可以看我之前的文章。

SSD代码我这里是在b站up主Bubbliiiing的pytorch版SSD的基础上进行的实现。

---

环境说明

windows10

cuda10.2

cudnn8.2.1

pytorch1.7

tensorrt8.2.5.1

python 3.7

显卡：NVIDIA 1650 4G(比较拉跨)

注：linux下我还没有试，可能有些代码需要改，而且trt的版本也会受影响

---

首先说一下网络输出部分我修改了哪些。

【其实这部分不看也可以，可以直接跳过1，2，3节，直接从转onnx开始看】

再转onnx之前需要修改一些地方，我们知道SSD有6个输出，代码中的输出output是包含了loc【位置】，conf【每个先验框对应类的置信度】,先验框数量。在代码中分train模式和test模式【这两个模式好像在后面的版本中进行了整合】。

train模式下的输出为：

output = (
                 loc.view(loc.size(0), -1, 4),
                 conf.view(conf.size(0), -1, self.num_classes),
                 self.priors
             )

test模式下的输出为：

        if self.phase == "test":
            output = self.detect.forward(
                loc.view(loc.size(0), -1, 4),  # (batch_size, num_anchors,4) = (batch_size, num_anchors,(x,y,w,h))
                self.softmax(conf.view(conf.size(0), -1, self.num_classes)),  # (batch_size,num_anchors, num_classes)
                self.priors              
            )

因此我们首先要做的是修改一些输出，至于为什么要修改呢，是因为如果我导ONNX的时候进行检测的时候发现没有任何检测结果，分析内部数据的时候发现在num_classes以及对应的conf维度上无结果，因此当我以train模式导出onnx模式并在外面采用detect时发现是可以的。【反正就是我的个人经验得来的，大家用就可以了】

目录

1.修改ouputput

2.修改box解码

3.获得输出

4.转onnx[可以直接从这部分看]

5.ONNX推理

6.engine推理

 7.FPS测试

---

1.修改ouputput

所以需要将output改为以下【在nets/ssd.py】：

        if self.phase == "test":
            output = self.detect.forward(
                loc.view(loc.size(0), -1, 4),  # (batch_size, num_anchors,4) = (batch_size, num_anchors,(x,y,w,h))
                self.softmax(conf.view(conf.size(0), -1, self.num_classes)),  # (batch_size,num_anchors, num_classes)
                self.priors              
            )
        else:
            # output = (
            #     loc.view(loc.size(0), -1, 4),
            #     conf.view(conf.size(0), -1, self.num_classes),
            #     self.priors
            # )
            loc = torch.tensor(loc.view(loc.size(0), -1, 4)),
            conf = torch.tensor(conf.view(conf.size(0), -1, self.num_classes)),
            self.priors = torch.tensor(self.priors)
        #return output
        return loc, conf, self.priors

2.修改box解码

接下来再去nets/ssd_layers.py中的Dectect中的forward()中添加一行：

prior_data = prior_data.cpu()

也就是把先验框放在cpu上，不然在box解码的时候会出问题。

3.获得输出

然后在工程文件中的ssd.py中的检测部分添加下面的代码：

因为我们的输出现在有三个，所以要将其送入detect中【如果你是pytorch1.5以后版本要加上forward不然会报错】

            preds = self.net(photo)
            if self.onnx:
                loc = preds[0]
                conf = preds[1]
                priors = preds[2]
                preds = self.detect.forward(loc, nn.Softmax(dim=-1).forward(conf), priors)

4.转onnx

运行以下代码，仅需要修改torch权重路径即可以及类别：

python torch2onnx.py

torch转onnx模型我这里以及写好了，如果你是自己的数据集，需要修改num_classes，ckpt中torch的模型，output_namse和input_names是输出以及输入结点，因为有三个输出，所以是三个结点名字。最终导出的模型会保存在model_data文件下。

这里导出onnx有两个模式，你可以选择是否开启simplity，如果开启改功能可以更详细的看清楚每个卷积输出的尺寸大小，而且也会对onnx模型进行一定的优化。

import onnx
from nets.ssd import get_ssd
import torch
from utils.config import Config
from onnxsim import simplify
import numpy as np

Simplity = False

output_path = "model_data/ssd.onnx"
num_classes = 21
model = get_ssd('train', num_classes)
model_dict = model.state_dict()
device = torch.device('cuda')
ckpt = torch.load('./model_data/ssd_weights.pth',map_location=device)
ckpt = {k: v for k, v in ckpt.items() if np.shape(model_dict[k]) == np.shape(ckpt[k])}
model_dict.update(ckpt)
model.load_state_dict(model_dict)
model.eval()
model.to(device)
x = torch.zeros(1, 3, 300, 300).to(device)
output_names = ["output0","output1", "output2"]
input_names = ["images"]
torch.onnx.export(model, x, output_path, verbose=True, input_names=input_names,
                  output_names=output_names, do_constant_folding=True, opset_version=12)
if Simplity:
    onnx_model = onnx.load(output_path)  # load onnx model
    model_simp, check = simplify(onnx_model)
    assert check, "Simplified ONNX model could not be validated"
    onnx.save(model_simp, output_path)
    print('finished exporting onnx')

输出onnx部分图 

5.ONNX推理

在推理前，你需要进入ssd.py中修改以下部分，将ONNX设置为True：

_defaults = {
    "model_path"    : 'model_data/ssd.onnx', # 权重路径
    "classes_path"  : 'model_data/voc_classes.txt',
    "confidence"    : 0.5,
    "nms_iou"       : 0.45,
    "cuda"          : True,
}

#---------------------------------------------------#
#   初始化SSD
#---------------------------------------------------#
def __init__(self, input_shape=None, ONNX=False, TRT=False, **kwargs):
    if input_shape is None:
        input_shape = [300, 300]  # 支持300和512大小
    self.input_shape = input_shape
    self.onnx = ONNX  # 是否开启onnx推理
    self.engine = TRT  # 是否开启trt推理
    self.__dict__.update(self._defaults)
    self.class_names = self._get_class()
    self.generate()

注意：输入大小以及要和你onnx、engine输入大小一致！【我这里没做动态输入】 

然后运行predict.py即可。 

 ONNX推理这部分代码是参考了YOLOV5中的方式，代码如下：

import torch
import torch.nn as nn
import numpy as np

class DetectMultiBackend(nn.Module):
    def __init__(self, weights, device=torch.device('cpu'), fp16=False):
        super(DetectMultiBackend, self).__init__()
        cuda = torch.cuda.is_available()
        if weights.split('.')[-1] == "onnx":
            import onnxruntime
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
            session = onnxruntime.InferenceSession(weights, None)
            output_names = [x.name for x in session.get_outputs()]
            print("Output_names: ", output_names)
            meta = session.get_modelmeta().custom_metadata_map  # metadata
        self.__dict__.update(locals())
    def forward(self, im):
        global y
        if self.weights.split('.')[-1] == 'onnx':
            im = im.cpu().numpy()
            y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
        if isinstance(y, (list, tuple)):  # 多输出
            return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
        else:
            return self.from_numpy(y)

    def from_numpy(self, x):
        return torch.from_numpy(x).to(self.device) if isinstance(x, np.ndarray) else x

 

 推理结果

---

6.engine推理

将ssd.py中的model_path路径设置为engine路径，如下：

_defaults = {
    "model_path"    : 'model_data/ssd.engine',
    "classes_path"  : 'model_data/voc_classes.txt',
    "confidence"    : 0.5,
    "nms_iou"       : 0.45,
    "cuda"          : True,
}

将TRT功能打开，设置为True，同时注意网络输入尺寸:

def __init__(self, input_shape=None, ONNX=False, TRT=True, **kwargs):
    if input_shape is None:
        input_shape = [300, 300]
    self.input_shape = input_shape
    self.onnx = ONNX
    self.engine = TRT
    self.__dict__.update(self._defaults)
    self.class_names = self._get_class()
    self.generate()

运行predict.py，输入图像路径进行推理：

10/26/2022-17:38:05] [TRT] [I] [MemUsageChange] Init CUDA: CPU +421, GPU +0, now: CPU 5526, GPU 896 (MiB)
[10/26/2022-17:38:07] [TRT] [I] Loaded engine size: 131 MiB
[10/26/2022-17:38:09] [TRT] [I] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +326, GPU +70, now: CPU 5998, GPU 1097 (MiB)
[10/26/2022-17:38:10] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +108, GPU +88, now: CPU 6106, GPU 1185 (MiB)
[10/26/2022-17:38:10] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in engine deserialization: CPU +0, GPU +131, now: CPU 0, GPU 131 (MiB)
[10/26/2022-17:38:10] [TRT] [I] [MemUsageChange] Init cuBLAS/cuBLASLt: CPU +0, GPU +10, now: CPU 5983, GPU 1177 (MiB)
[10/26/2022-17:38:10] [TRT] [I] [MemUsageChange] Init cuDNN: CPU +0, GPU +8, now: CPU 5983, GPU 1185 (MiB)
[10/26/2022-17:38:10] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +168, now: CPU 0, GPU 299 (MiB)
model_data/ssd.engine model, anchors, and classes loaded.
Input image filename:

如果是视频推理就运行video.py即可

---

 7.FPS测试

如果是用torch模型进行测试，将FPS_test.py中的TORCH设置为True，然后在ssd.py中输入权重路径，然后运行FPS_test.py即可。

如果是用onnx或者engine模型进行测试，将FPS_test.py中的TORCH设置为False，然后在ssd.py中输入权重路径，然后运行FPS_test.py即可。

如果输入大小为300 * 300，在我的显卡上FPS如下，可以看到其实提升还是可以的，提升了十几帧：

NVIDIA 1650 4G
# engine:300*300 41FPS
# onnx:300*300 4FPS
# torch:300*300 28FPS

---

 详细使用请看readme.md

代码： 

https://github.com/YINYIPENG-EN/SSD_tensorRT_pytorch.git

权重：

链接：https://pan.baidu.com/s/1YAF_cruDG254ZZD5InPsbg 
提取码：yypn