From 31b3e7141d2acb0cb1ff298d2b7fca3911889f25 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Wed, 2 Sep 2020 21:18:28 +0200 Subject: [PATCH] Initial commit. --- speed.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100755 speed.py diff --git a/speed.py b/speed.py new file mode 100755 index 0000000..f682a16 --- /dev/null +++ b/speed.py @@ -0,0 +1,28 @@ +#!/usr/bin/env python + +import time, torch + +if torch.cuda.is_available(): + device = torch.device('cuda') +else: + device = torch.device('cpu') + +nb_runs = 10000 +d1, d2, d3 = 50000, 256, 512 + +a, b = torch.rand(d1, d2).to(device), torch.rand(d2, d3).to(device) + +start_time = time.perf_counter() +for k in range(nb_runs): + c = a @ b +duration = time.perf_counter() - start_time + +nb_flop = float(nb_runs * d1 * d2 * d3) +speed = nb_flop / duration + +for u in [ '', 'K', 'M', 'G', 'T', 'P' ]: + if speed < 1e3: break + speed /= 1e3 + +print(f'{speed:.02f} {u}flops on {device}') + -- 2.20.1