【机器学习】使用NVIDIA GPU训练模型，CUDA对应版本torch下载安装避坑指南

Albert_6

于 2024-09-04 11:41:33 发布

阅读量220

点赞数 6

分类专栏：人工智能机器学习测试开发文章标签：机器学习人工智能

本文链接：https://blog.csdn.net/Albert_6/article/details/140526290

版权

5 篇文章 0 订阅

订阅专栏

2 篇文章 0 订阅

订阅专栏

1 篇文章 0 订阅

订阅专栏

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()

关注

专栏目录