Yolov5 适配TensorRT的问题记录

2023年10月09日09:44:21

author：Edge

记录：在适配 agx 系列设备中，使用 Trt调用 gpu推理的过程中，出现了一个问题：

File "forward_trt.py", line 176, in res_pt = model_pt.run(img) File "forward_trt.py", line 119, in run masks = process_mask(proto[i], det[:, 6:], det[:, :4], im.shape[2:], upsample=True,foreground_threshold = self.fmask_threshold) # HWC File "xxx/general.py", line 54, in process_mask c, mh, mw = protos.shape # CHW ValueError: not enough values to unpack (expected 3, got 2)

这个错误意味着，你的输出 proto 并不是正确的 proto。

网络中，带有五个输出，如下图：

其中第一个名为 output 的输出，是检测结果的输出，维度是[1,25200,54]，另一个 name 为 462 的输出是分割分支的输出，维度是[1,32,160,160]。

作者写的 forward 中：

x
elif self.engine:  # TensorRT
            if self.dynamic and im.shape != self.bindings['images'].shape:
                i = self.model.get_binding_index('images')
                self.context.set_binding_shape(i, im.shape)  # reshape if dynamic
                self.bindings['images'] = self.bindings['images']._replace(shape=im.shape)
                for name in self.output_names:
                    i = self.model.get_binding_index(name)
                    self.bindings[name].data.resize_(tuple(self.context.get_binding_shape(i)))
            s = self.bindings['images'].shape
            assert im.shape == s, f"input size {im.shape} {'>' if self.dynamic else 'not equal to'} max model size {s}"
            self.binding_addrs['images'] = int(im.data_ptr())
            self.context.execute_v2(list(self.binding_addrs.values()))
            y = [self.bindings[x].data for x in sorted(self.output_names)]

y = [self.bindings[x].data for x in sorted(self.output_names)]这里获取所有 name 的输出，与 predict 中匹配不上，导致的错误，改成：

y = [self.bindings[x].data for x in ['output', '462']]

就可以正常推理输出了，具体可以看 我提交的PR：

https://github.com/ultralytics/yolov5/pull/12209