search

laugh12321 / TensorRT-YOLO

🚀 你的YOLO部署神器。TensorRT Plugin、CUDA Kernel、CUDA Graphs三管齐下，享受闪电般的推理速度。| Your YOLO Deployment Powerhouse. With the synergy of TensorRT Plugins, CUDA Kernels, and CUDA Graphs, experience lightning-fast inference speeds.
https://github.com/laugh12321/TensorRT-YOLO
GNU General Public License v3.0
540 stars 67 forks source link

[Help]: Key Error When Inferencing In Jupyter #26

Closed Tom-Teamo closed 4 months ago

Tom-Teamo commented 4 months ago

由于我想调试代码以便更好地了解一下YOLO的具体流程，但是我在python Jupyter使用如下代码：

engine_path = '/home/TensorRT-YOLO/demo/detect/models/yolov8s.engine'
model = TRTYOLO(engine_path)
image_path = '/home/TensorRT-YOLO/demo/detect/images'
batcher = ImageBatcher(input_path=image_path, batch_size=model.batch_size, imgsz=model.imgsz, dtype=model.dtype, dynamic=model.dynamic)

model.warmup()
for batch, images, batch_shape in batcher:
    detections = model.infer(batch, batch_shape)

会在post_process中出现KeyError: 'num_dets'，在post_process函数中，outputs.keys()如下，所以会出现KeyError

dict_keys(['num_detections', 'detection_boxes', 'detection_scores', 'detection_classes'])

但是在python CLI中：trtyolo infer -e models/yolov8s.engine -i images -o output -l labels.txt 并没有报错

两者的源代码版本是一致的，请问这是为什么？

laugh12321 commented 4 months ago

重新导出模型即可。

新版本导出的模型对节点名称做了修改，参考https://github.com/laugh12321/TensorRT-YOLO/commit/2d3d630625430e91bd896c2cf5abd26aa7a9b041

你使用了老版本导出的模型，但是用新版本代码进行推理。