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显存测试-StubbornHuang Blog纵使晴明无雨色,入云深处亦沾衣。