mxnet_numpy_performance_test.py

import numpy as np
from time import time

size = 4096
N = 20
def timed_matrix_multiply(A, B):
    t = time()
    for i in range(N):
        np.dot(A, B)
    elapsed = time() - t
    return elapsed
    
A, B = np.random.uniform(size=(size, size)), np.random.uniform(size=(size, size))
elapsed = timed_matrix_multiply(A, B)
print('NumPy : Dotted two %dx%d matrices in %0.2f s.' % (size, size, elapsed / N))

import mxnet.numpy as np
from mxnet import npx 
npx.set_np()

def timed_matrix_multiply(A, B):
    t = time()
    for i in range(N):
        C = np.dot(A, B)
        C.wait_to_read()
    elapsed = time() - t
    return elapsed

A, B = np.random.uniform(size=(size, size)), np.random.uniform(size=(size, size))
elapsed = timed_matrix_multiply(A, B)
print('mxnet.numpy : Dotted two %dx%d matrices in %0.2f s.' % (size, size, elapsed / N))

A, B = np.random.uniform(size=(size, size), ctx=npx.gpu()), np.random.uniform(size=(size, size), ctx=npx.gpu())
elapsed = timed_matrix_multiply(A, B)
print('mxnet.numpy on GPU : Dotted two %dx%d matrices in %0.2f s.' % (size, size, elapsed / N))