GeForce 驱动程序下载
https://www.nvidia.cn/geforce/drivers/
选择与自己设备硬件匹配的版本,下载安装。
查看CUDA驱动版本
NVIDIA - 系统信息 - 组件,在3D设置中可以看到cuda版本信息,
这里大版本号是11.3。
查询CUDA版本是否与GPU驱动程序兼容
https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html
CUDA Toolkit 下载
https://developer.nvidia.com/
输入后敲回车搜索:
PyTorch官网查找以前的版本
https://pytorch.org/get-started/previous-versions/
找到使用pip命令安装的
不要直接使用这个命令安装,因为官网在线安装巨慢,而且会连接超时中断
torch版本是 1.12.1+cu113
可以去阿里云开源镜像站 https://mirrors.aliyun.com/pytorch-wheels/
搜索下载,本地安装。注意对应的python版本
下载之后,使用pip安装,路径改成你的路径
pip install C:\Users\Albert\Downloads\torch-1.12.1+cu113-cp38-cp38-win_amd64.whl
torchvision和torchaudio比较小,可以带上对应的版本号直接pip在线安装
pip install torchvision==0.13.1
pip install torchaudio==0.12.1
使用torch.cuda.is_available()验证,返回true即可用
如下定义设备,兼容性更好:
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
加载模型时,就可以指定映射到cuda设备上(即GPU)
model = NeuralNetWork().to(device)
model.load_state_dict(torch.load(f"./model_{CODE_TYPE}.pth", map_location=torch.device("cuda")))
model.eval()