在修改代码以访问 YOLOv5 for CPU 时需要帮助

Question

我正在尝试使用 Pytorch、OpenCV 和 YOLOv5 实现一个对象检测程序，以检测 YouTube 视频中的对象和对象类型。但是，在运行时，输出控制台显示程序尝试运行的 YOLO 版本是针对 CUDA 的。我希望使用 YOLO for CPU 来实现该项目，因为我的计算机没有为 CUDA 设置。我寻求有关如何修改程序以将 YOLO 用于 CPU 的帮助？非常感谢！

GitHub Gist 链接

https://gist.github.com/neitherjames/c3b91033eca3794f8170ee51ee5357d4

Answer 1

在您的__init__ 中，它表明您为系统提供了选择 CUDA（如果可用）的选项。您可以通过在第 23 行声明 self.device = 'cpu' 来强制它在 cpu 上运行。然后在第 48 行调用 self.model.to(self.device) 时，将模型发送到 cpu。