pytorch版本与torchvision版本不匹配问题处理
- 问题
- 问题复现
- 解决方法
- 两点注意内容
- 其一:pytorch版本与torchvision版本对应关系
- 其二:CPU版本或GPU版本问题
问题
在新环境中,利用yolov8训练模型的时候报错,错误内容如下:
NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'torchvision::nms' is only available for these backends: [CPU, Meta, QuantizedCPU, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMeta, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, BatchedNestedTensor, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PreDispatch, PythonDispatcher].
这个问题出现出现的原因是因为pytorch版本与torchvision版本不匹配导致利用pip查看安装的包,果然存在问题,发现pytorch版本为GPU版本,但是torchvision为CPU版本。
问题复现
我问题出现的原因是在新环境中,利用如下安装命令:
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
安装pytorch的时候,pytorch由于很大,下载很慢,所以我单独下载了pytorch的whl包进行了安装,然后利用如下命令安装了torchvision和torchaudio:
pip3 install torchvision torchaudio
正是这句安装命令导致了问题,这句安装命令会默认安装torchvision的CPU版本。
解决方法
解决方法很简单,卸载掉torchvision,然后安装GPU版本的torchvision即可。
安装命令如下:
pip3 install torchvision --index-url https://download.pytorch.org/whl/cu118
两点注意内容
其一:pytorch版本与torchvision版本对应关系
pytorch版本与torchvision版本对应关系可以参考官网,链接地址如下:
https://github.com/pytorch/vision
下翻可找到对应对应关系表,如下所示:
其二:CPU版本或GPU版本问题
注意,你安装pytorch为CPU版本,那么torchvision也需要是CPU版本,同样,pytorch为GPU版本,那么torchvision也需要是GPU版本。