出现以下错误提示
Traceback (most recent call last):
File "train.py", line 470, in <module>
train(hyp, opt, device, tb_writer)
File "train.py", line 333, in train
plots=epoch == 0 or final_epoch) # plot first and last
File "/project/train/src_repo/yolov5-master_jishi/test.py", line 193, in test
plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions
File "/project/train/src_repo/yolov5-master_jishi/utils/general.py", line 951, in output_to_target
return np.array(targets)
File "/usr/local/lib/python3.7/dist-packages/torch/_tensor.py", line 678, in __array__
return self.numpy()
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.
我的解决办法是将utils/general.py文件中
def output_to_target(output, width, height):
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
if isinstance(output, torch.Tensor):
output = output.cpu().numpy()
targets = []
for i, o in enumerate(output):
if o is not None:
for pred in o:
box = pred[:4]
w = (box[2] - box[0]) / width
h = (box[3] - box[1]) / height
x = box[0] / width + w / 2
y = box[1] / height + h / 2
conf = pred[4]
cls = int(pred[5])
# 列表里边有tensor数据,需要.cpu()
targets.append([i, cls, x.cpu(), y.cpu(), w.cpu(), h.cpu(), conf.cpu()])
return np.array(targets)
targets.append([i, cls, x.cpu(), y.cpu(), w.cpu(), h.cpu(), conf.cpu()])
x,y,w,h,conf 这几个tensor全部加上.cpu(),就可以跑通代码了。