Skip to content

Instantly share code, notes, and snippets.

@entron
Created August 14, 2017 18:44
Show Gist options
  • Save entron/7ce80ecc1eef684f93817f99118a5093 to your computer and use it in GitHub Desktop.
Save entron/7ce80ecc1eef684f93817f99118a5093 to your computer and use it in GitHub Desktop.
Test pytorch linear algebra speed
import torch
import time
# --- Test 1
N = 10
n = 1000
A = torch.randn(n, n)
B = torch.randn(n, n)
t = time.time()
for i in range(N):
C = torch.mm(A, B)
td = time.time() - t
print("dotted two (%d,%d) matrices in %0.1f ms" % (n, n, 1e3 * td / N))
# --- Test 2
N = 100
n = 4000
A = torch.randn(n)
B = torch.randn(n)
t = time.time()
for i in range(N):
C = torch.dot(A, B)
td = time.time() - t
print("dotted two (%d) vectors in %0.2f us" % (n, 1e6 * td / N))
# --- Test 3
m, n = (2000, 1000)
A = torch.randn(m, n)
t = time.time()
[U, s, V] = torch.svd(A, some=False)
td = time.time() - t
print("SVD of (%d,%d) matrix in %0.3f s" % (m, n, td))
# --- Test 4
n = 1500
A = torch.randn(n, n)
t = time.time()
w, v = torch.eig(A)
td = time.time() - t
print("Eigendecomp of (%d,%d) matrix in %0.3f s" % (n, n, td))
# --- Test 5
N = 100
n = 2000
A = torch.randn(n, n)
t = time.time()
for i in range(N):
B = torch.exp(A)
td = time.time() - t
print("Element wise exp of (%d,%d) matrix in %0.1f ms" % (n, n, 1e3 * td / N))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment