yolov9目标检测报错AttributeError: ‘list‘ object has no attribute ‘device‘

深度学习

文章目录

深度学习
前言

前言

yolov9运行自己训练的模型时，出现以下错误：

bash 复制代码

root@b219ae83c78f:/yolov9# python detect.py --source './data/images/horses.jpg' --img 640 --device 0 --weights runs/train/yolov9-c8/weights/best.pt --name yolov9_c_c_640_detect2
detect: weights=['runs/train/yolov9-c8/weights/best.pt'], source=./data/images/horses.jpg, data=data/coco128.yaml, imgsz=[640, 640], conf_thres=0.25, iou_thres=0.45, max_det=1000, device=0, view_img=False, save_txt=False, save_conf=False, save_crop=False, nosave=False, classes=None, agnostic_nms=False, augment=False, visualize=False, update=False, project=runs/detect, name=yolov9_c_c_640_detect2, exist_ok=False, line_thickness=3, hide_labels=False, hide_conf=False, half=False, dnn=False, vid_stride=1
YOLO 🚀 v0.1-104-g5b1ea9a Python-3.8.12 torch-1.11.0a0+b6df043 CUDA:0 (NVIDIA TITAN V, 12057MiB)

Fusing layers... 
yolov9-c summary: 604 layers, 50880768 parameters, 0 gradients, 237.6 GFLOPs
Traceback (most recent call last):
  File "detect.py", line 231, in <module>
    main(opt)
  File "detect.py", line 226, in main
    run(**vars(opt))
  File "/opt/conda/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "detect.py", line 102, in run
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
  File "/yolov9/utils/general.py", line 905, in non_max_suppression
    device = prediction.device
AttributeError: 'list' object has no attribute 'device'

File "/yolov9/utils/general.py", line 905, in non_max_suppression

这行代码错误，应该是照抄了yolov5的代码

bash 复制代码

 if isinstance(prediction, (list, tuple)):  # YOLO model in validation model, output = (inference_out, loss_out)
        prediction = prediction[0]  # select only inference output

    device = prediction.device
    mps = 'mps' in device.type  # Apple MPS
    if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
        prediction = prediction.cpu()
    bs = prediction.shape[0]  # batch size
    nc = prediction.shape[1] - nm - 4  # number of classes
    mi = 4 + nc  # mask start index
    xc = prediction[:, 4:mi].amax(1) > conf_thres  # candidates

改成以下代码，问题解决，

bash 复制代码

    if isinstance(prediction, (list, tuple)):  # YOLO model in validation model, output = (inference_out, loss_out)
        processed_predictions = []
        for pred_tensor in prediction:
            processed_tensor = pred_tensor[0]
            processed_predictions.append(processed_tensor)
        #prediction = prediction[0]  # select only inference output
        prediction = processed_predictions[0]

    device = prediction.device
    mps = 'mps' in device.type  # Apple MPS
    if mps:  # MPS not fully supported yet, convert tensors to CPU before NMS
        prediction = prediction.cpu()
    bs = prediction.shape[0]  # batch size
    nc = prediction.shape[1] - nm - 4  # number of classes
    mi = 4 + nc  # mask start index
    xc = prediction[:, 4:mi].amax(1) > conf_thres  # candidates

完美解决。