0. 安装Nvidia驱动
ubuntu24.04的安装非常简单,在安装界面,选择为"图形化和其他硬件安装驱动",重启后即有原版Nvidia驱动(如图Nvidia X xxx)
1.确定电脑上是否有NvidiaGPU且安装好Nvidia驱动
import torch
print(torch.version.cuda)
cuda简介
CUDA是一种并行计算平台和应用程序编程接口(API),用于利用NVIDIA图形处理器(GPU)进行通用计算。它使开发人员能够使用GPU进行并行计算,通过将计算任务分发到多个GPU核心上同时执行,从而加速计算过程。
CUDA支持多种编程语言,包括C、C++、Fortran和Python,使开发人员能够在GPU上编写高性能的并行计算应用程序。
理论上,也能用C++访问cuda API。
2. 加载tensor到Nvidia GPU
a = torch.tensor(range(5))
if torch.cuda.is_available():
a = a.to('cuda:0') #等价
a = a.to('cuda')
else:
a = a.to('cpu')
print(a.device)