测试Torch gpu 是否可用

1
2
3
import torch

print(torch.version.cuda)
1
2
3
4
import 	torch
import time
print(torch.__version__)
print(torch.cuda.is_available())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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))