终于更新了，本篇是实现了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 outputreturn 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 npSimplity = Falseoutput_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 modelmodel_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_shapeself.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 npclass 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 onnxruntimeproviders = ['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  # metadataself.__dict__.update(locals())def forward(self, im):global yif 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_shapeself.onnx = ONNXself.engine = TRTself.__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

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

解析目标检测全流程！附代码数据
↑↑↑关注后"星标"Datawhale 每日干货 & 每月组队学习,不错过 Datawhale干货作者:王程伟,算法工程师,Datawhale成员在计算机视觉中,红外弱 ...
神经网络: SSD目标检测网络解析
一.SSD(Single Shot MultiBox Detector)简介 SSD是2016年ICCV的一篇论文,是目前为止主要的目标检测算法. 算法的主要优点: 速度比Faster-Rcnn快,精 ...
yolov3网络结构图_目标检测——YOLO V3简介及代码注释（附github代码——已跑通）...
GitHub: liuyuemaicha/PyTorch-YOLOv3github.com 注:该代码fork自eriklindernoren/PyTorch-YOLOv3,该代码相比master分 ...
丢弃Transformer！旷视和西安交大提出基于FCN的端到端目标检测网络
点击上方,选择星标或置顶,不定期资源大放送! 阅读大概需要15分钟 Follow小博主,每天更新前沿干货本文作者:王剑锋 | 编辑:Amusi https://zhuanlan.zhihu.com ...
retinanet 部署_RetinaNet: Focal loss在目标检测网络中的应用
介绍 RetinaNet是2018年Facebook AI团队在目标检测领域新的贡献.它的重要作者名单中Ross Girshick与Kaiming He赫然在列.来自Microsoft的Sun Jia ...
MobileViT: 一种更小，更快，高精度的轻量级Transformer端侧网络架构（附代码实现）...
点击上方,选择星标或置顶,不定期资源大放送! 阅读大概需要5分钟 Follow小博主,每天更新前沿干货 [导读]之前详细介绍了轻量级网络架构的开源项目,详情请看深度学习中的轻量级网络架构总结与代码实现 ...
Pytorch搭建SSD目标检测平台
学习前言什么是SSD目标检测算法源码下载 SSD实现思路一.预测部分 1.主干网络介绍 2.从特征获取预测结果 3.预测结果的解码 4.在原图上进行绘制二.训练部分 1.真实框的处理 2.利用 ...
计算机视觉算法——目标检测网络总结
计算机视觉算法--目标检测网络总结计算机视觉算法--目标检测网络总结 1. RCNN系列 1.1 RCNN 1.1.1 关键知识点--网络结构及特点 1.1.2 关键知识点--RCNN存在的问题 1 ...
收藏 | 如何定义目标检测网络的正负例：Anchor-based
点上方计算机视觉联盟获取更多干货仅作学术分享,不代表本公众号立场,侵权联系删除转载于:作者:知乎-Lighthouse 地址:https://www.zhihu.com/org/lighthous ...
《基于海思35xx nnie引擎进行经典目标检测算法模型推理》视频课程介绍
前言沉寂两个月,终于将新的视频课程<<基于海思35xx nnie引擎进行经典目标检测算法模型推理>>(其链接为https://edu.csdn.net/course/deta ...

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

1.修改ouputput

2.修改box解码

3.获得输出

4.转onnx

5.ONNX推理

6.engine推理

7.FPS测试

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

最新文章

热门文章