10a008f68dc6c3865fd3a0e00740016055260306
[pytorch.git] / speed.py
1 #!/usr/bin/env python
2
3 import time, torch
4
5 if torch.cuda.is_available():
6     device = torch.device('cuda')
7     sync = lambda: torch.cuda.synchronize()
8 else:
9     device = torch.device('cpu')
10     sync = lambda: None
11
12 nb_runs = 10000
13 d1, d2, d3 = 2048, 2048, 2048
14
15 a, b = torch.rand(d1, d2).to(device), torch.rand(d2, d3).to(device)
16
17 sync
18 start_time = time.perf_counter()
19 for k in range(nb_runs):
20     c = a @ b
21 sync
22 duration = time.perf_counter() - start_time
23
24 nb_flop = float(nb_runs * d1 * d2 * d3 * 2) # 1 multiply-and-add is 2 ops
25 speed = nb_flop / duration
26
27 for u in [ '', 'K', 'M', 'G', 'T', 'P' ]:
28     if speed < 1e3: break
29     speed /= 1e3
30
31 print(f'{speed:.02f} {u}flops on {device}')
32