Pytorch的CUDA版本安装使用教程

程序员miki2025-09-05 21:14

CUDA

什么是 CUDA

CUDA 是英伟达公司开发的显卡的人工智能算法驱动程序，原本我们训练模型是基于的CPU，但是我们很清楚，单靠CPU和内存来训练模型是远远达不到要求的，而且训练的速度会比较的慢，因此我们就需要我们的显卡来训练我们的模型。

限制条件

为了使用 CUDA，我们需要一个被 CUDA 支持的英伟达的显卡

官网查看那些显卡可以使用CUDA GPU Compute Capability | NVIDIA Developer

CUDA安装

CUDA驱动安装

但是大部分情况下，一般是不需要单独安装 cuda 计算驱动下载地址

选择你的系统，比如Windows。

Pytorch的CUDA版本

还有另外的下载方式，比如我们使用Pytorch来进行模型的训练，Pytorch官方也有对应的CUDA版本，直接下载CUDA版本的Pytorch。

Pytorch官网

在下载的时候，记得删除旧版本的Pytorch，这里我就不赘述了。

CUDA使用

CUDA是否存活

python 复制代码

print(torch.cuda.is_available())

CUDA是否可用，创立一个设备对象

python 复制代码

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

将模型丢入到CUDA中

python 复制代码

# 1. 创建设备对象
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. 将模型放到设备上
model.to(device)

# 3. 将数据放到设备上
inputs, labels = inputs.to(device), labels.to(device)