众所周知啊,在我们安装 GPU 版的 pytorch 之前,需要安装 CUDA,CUDA 的安装教程在此不过多追叙,详情可以看官方文档。
咱首先要做的是确认我们的 CUDA 的版本
知道 CUDA 版本之后,去pytorch官网一键下载就好了
查看 torch 以及 torchvision 的版本
import torch
import torchvision
print("pytorch version:",torch.__version__)
print("torchvision version:",torchvision.__version__)
注:这里如果你CUDA版本是11.6,并按照我上边的流程去弄,接下来会报一个错
如果 CUDA 是其它版本的可以忽略我下边说的直接点赞加收藏就好了【】
“NotImplementedError: Could not run ‘torchvision::nms‘ with arguments from the ‘CUDA‘ backend.”
这里解释一下为什么会报错,原因是你 CUDA11.6 在官方网首页直接下载的 pytorch 三宝合集包里边,torchvision 这个是默认给你下 cpu 版本的,所以这边咱需要换一下它的版本
创造亚当之兄弟抽支烟、点个赞