yolov5 导出torchscript报错:RuntimeError: Input type (torch.cuda.FloatTensor) and weight type

github ultralytics的yolov5 v6.0版本代码导出torchscript导出成功了但是detect.py推理的时候报RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same

这里有两种解决思路

1. 将输入input和weight统一在一个设备上(我采用的这个,比较简单粗暴)

排查代码可以发现,在detect.py中有个DetectMultiBanckend这是用于载入不同类型的模型的类,该类定义在 根目录/models/common.py中,查看该函数对应的 jit(即torchscript) 处代码

 可见这里少了一个参数,map_location=device,加上即可

2. 参考(97条消息) Pytorch避坑之:RuntimeError: Input type(torch.cuda.FloatTensor) and weight type(torch.FloatTensor) shoul_暖仔会飞的博客-CSDN博客

猜你喜欢

转载自blog.csdn.net/m0_52948781/article/details/124671318