Pytorch – 检测CUDA、cuDNN以及GPU版本的Pytorch是否安装成功、GPU显存测试
1 测试代码
import time
import torch
if __name__ == '__main__':
print('torch版本：'+torch.__version__)
print('cuda是否可用：'+str(torch.cuda.is_available()))
print('cuda版本：'+str(torch.version.cuda))
print('cuda数量:'+str(torch.cuda.device_count()))
print('GPU名称：'+str(torch.cuda.get_device_name()))
print('当前设备索引：'+str(torch.cuda.current_device()))
device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
print(device)
print(torch.rand(3, 3).cuda())
for i in range(1,100000):
start = time.time()
a = torch.FloatTensor(i*100,1000,1000)
a = a.cuda() #a = a
a = torch.matmul(a,a)
end = time.time() - start
print(end)
结果：
torch版本：1.7.1
cuda是否可用：True
cuda版本：11.0
cuda数量:1
GPU名称：NVIDIA GeForce RTX 3090
当前设备索引：0
cuda:0
tensor([[0.3954, 0.7494, 0.8901],
[0.6576, 0.1618, 0.6446],
[0.1332, 0.2121, 0.8951]], device='cuda:0')
如果显示cuda可用，并且pytorch tensor可正常放到cuda上面去计算，那么就表明整个pytorch+cuda环境搭建成功，后面的for循环主要是为了测试GPU的显存是否能正常被pytorch调用。
欢迎扫码关注我的微信公众号，及时获取文章更新
本文作者：StubbornHuang
版权声明：本文为站长原创文章，如果转载请注明原文链接！
原文标题：Pytorch – 检测CUDA、cuDNN以及GPU版本的Pytorch是否安装成功、GPU显存测试
原文链接：https://www.stubbornhuang.com/1500/
发布于：2021年08月06日 10:33:26
修改于：2023年06月26日 21:23:06
声明：本站所有文章，如无特殊说明或标注，均为本站原创发布。任何个人或组织，在未征得本站同意时，禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益，可联系我们进行处理。
文章末尾
评论
72