终于更新了，本篇是实现了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
所以需要将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
