From: Francois Fleuret <francois@fleuret.org>
Date: Wed, 2 Sep 2020 19:18:28 +0000 (+0200)
Subject: Initial commit.
X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=31b3e7141d2acb0cb1ff298d2b7fca3911889f25;p=pytorch.git

Initial commit.
---

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}')
+