在我们跑深度学习的时候,绕不开的一定是gpu显卡,但是如何检测是否使用gpu那?下面我讲提供一段代码,可以检测自己是否开启了gpu或含有gpu。
import torch# 检查是否支持 CUDA(即是否有可用的 NVIDIA GPU)
flag = torch.cuda.is_available()
print(flag) # 输出 True 或 False,表示是否支持 CUDA# 设置设备(GPU 或 CPU)
ngpu = 1 # 假设使用 1 个 GPU
# 如果支持 CUDA 且 ngpu > 0,则使用 GPU(cuda:0 表示第一个 GPU),否则使用 CPU
device = torch.device("cuda:0" if torch.cuda.is_available() and ngpu > 0 else "cpu")
print(device) # 输出设备类型(cuda:0 或 cpu)# 如果使用 GPU,打印 GPU 的名称
print(torch.cuda.get_device_name(0)) # 输出第一个 GPU 的名称# 创建一个随机张量并将其移动到 GPU(如果使用 GPU)
print(torch.rand(3, 3).cuda()) # 创建一个 3x3 的随机张量并移动到 GPU(如果支持 CUDA)import torch # 这里其实可以不使用,只是我的习惯而已~# 获取 CUDA 的版本号
cuda_version = torch.version.cuda
print("Cuda Version:", cuda_version) # 输出 CUDA 的版本号# 获取 cuDNN 的版本号
cudnn_version = torch.backends.cudnn.version()
print("Cudnn Version:", cudnn_version) # 输出 cuDNN 的版本号