PyTorch GPU环境验证与CUDA信息查看

PyTorch GPU环境验证与CUDA信息查看

参考B站炮哥带你学

新建一个Python文件，输入以下代码来测试PyTorch环境：

importnumpyasnpimportpandasaspdimporttorch# 检查CUDA是否可用flag=torch.cuda.is_available()print(f"CUDA可用:{flag}")# 设置设备ngpu=1device=torch.device("cuda:0"ifflagelse'cpu')print(f"设备:{device}")# 打印GPU信息print(f"GPU名称:{torch.cuda.get_device_name(0)}")print(f"随机张量测试:{torch.rand(3,3).cuda()}")# 查看CUDA和cuDNN版本cuda_version=torch.version.cudaprint(f"CUDA版本:{cuda_version}")cudnn_version=torch.backends.cudnn.version()print(f"cuDNN版本:{cudnn_version}")