终于更新了，本篇是实现了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时发现是可以的。【反正就是我的个人经验得来的，大家用就可以了】
目录
4.转onnx[可以直接从这部分看]
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
转载：https://blog.csdn.net/z240626191s/article/details/127529298