Loading

深度学习:测试是否含有gpu的环境

在我们跑深度学习的时候,绕不开的一定是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 的版本号
posted @ 2025-02-10 10:03  虾饺爱下棋  阅读(131)  评论(0)    收藏  举报