问题分析
笔者在训练yolov5模型时出现如下问题:
Traceback (most recent call last):
File "train.py", line 460, in <module>
train(hyp, opt, device, tb_writer)
File "train.py", line 320, in train
plots=epoch == 0 or final_epoch) # plot first and last
File "D:\yolov5-3.1\test.py", line 197, in test
plot_images(img, output_to_target(output, width, height), paths, str(f), names) # predictions
File "D:\yolov5-3.1\utils\general.py", line 951, in output_to_target
return np.array(targets)
File "D:\Anaconda3\envs\yolov5-3\lib\site-packages\torch\tensor.py", line 621, 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.
解决问题
根据错误提示找到D:\Anaconda3\envs\yolov5-3\lib\site-packages\torch\tensor.py
文件的第621
行
将:
return self.numpy()
修改为:
return self.cpu().detach().numpy()
再次运行训练指令即可正常运行,问题解决