pytorch：测试GPU是否可用

import torchflag = torch.cuda.is_available()if flag:print("CUDA可使用")else:print("CUDA不可用")ngpu= 1# Decide which device we want to run ondevice = torch.device("cuda:0" if (torch.cuda.is...

文章共724字 · 阅读需要大约3分钟

一键AI生成摘要，助你高效阅读

问答

James_Bobo

49316人浏览 · 2019-11-18 09:19:16

James_Bobo · 2019-11-18 09:19:16 发布

import torch
flag = torch.cuda.is_available()
if flag:
    print("CUDA可使用")
else:
    print("CUDA不可用")

ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print("驱动为：",device)
print("GPU型号： ",torch.cuda.get_device_name(0))

在这里插入图片描述

import 	torch
import  time
print(torch.__version__)
print(torch.cuda.is_available())
# print('hello, world.')


a = torch.randn(10000, 1000)
b = torch.randn(1000, 2000)

t0 = time.time()
c = torch.matmul(a, b)
t1 = time.time()
print(a.device, t1 - t0, c.norm(2))

device = torch.device('cuda')
a = a.to(device)
b = b.to(device)

t0 = time.time()
c = torch.matmul(a, b)
t2 = time.time()
print(a.device, t2 - t0, c.norm(2))

t0 = time.time()
c = torch.matmul(a, b)
t2 = time.time()
print(a.device, t2 - t0, c.norm(2))