Блог

Я хочу использовать JAX для ускорения моего кода numpy на CPU, а затем на GPU. Вот мой пример кода, работающий на моем локальном компьютере (только CPU):

 import jax.numpy as jnp
from jax import random, jix
import numpy as np
import time

size = 3000

key = random.PRNGKey(0)
x =  random.normal(key, (size,size), dtype=jnp.float64)

start=time.time()
test = jnp.dot(x, x.T).block_until_ready()
print('Time of jnp: {}s'.format(time.time() - start))

x2=np.random.normal((size,size))

start=time.time()
test2 = np.dot(x2, x2.T)
print('Time of np: {}s'.format(time.time() - start))
  

Я получил предупреждение, и временные затраты следующие:

 /.../lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: 
UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Time: 0.45157814025878906s
Time: 0.005244255065917969s
  

Я сделал что-то не так здесь? Должен ли JAX также ускорять код numpy на процессорах?

Вероятно, существуют различия в производительности между Jax и Numpy, но в исходном сообщении разница во времени в основном сводится к ошибке при создании массива. Массив, используемый Jax, имеет форму 3000×3000, тогда как массив, используемый Numpy, представляет собой одномерный массив длиной 2. Первым аргументом numpy.random.normal является loc (т. Е. Среднее значение гауссова, из которого выполняется выборка). Аргумент ключевого слова size= следует использовать для указания формы массива.

 numpy.random.normal(loc=0.0, scale=1.0, size=None)
  

После внесения этого изменения производительность между Jax и Numpy меньше отличается.

 import time
import jax
import jax.numpy as jnp
import numpy as np

size = 3000

key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (size, size), dtype=jnp.float64)

start = time.time()
test = jnp.dot(x, x.T).block_until_ready()
print("Time of jnp: {:0.4f} s".format(time.time() - start))

x2 = np.random.normal(size=(size, size)).astype(np.float64)

start = time.time()
test2 = np.dot(x2, x2.T)
print("Time of np: {:0.4f} s".format(time.time() - start))
  

Результат одного запуска равен

 Time of jnp: 2.3315 s
Time of np: 2.8811 s
  

При измерении временной производительности следует собирать несколько запусков, потому что производительность функции — это разброс во времени, а не одно значение. Это можно сделать с помощью timeit.timeit функции стандартной библиотеки Python или %timeit магии в IPython и Jupyter Notebook.

 import time
import jax
import jax.numpy as jnp
import numpy as np

size = 3000

key = jax.random.PRNGKey(0)
xjnp = jax.random.normal(key, shape=(size, size), dtype=jnp.float64)
xnp = np.random.normal(size=(size, size)).astype(np.float64)

%timeit jnp.dot(xjnp, xjnp.T).block_until_ready()
# 2.03 s ± 39.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.dot(xnp, xnp.T)
# 3.41 s ± 501 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

xjnp = xjnp.astype(jnp.float32)
xnp = xnp.astype(np.float32)

%timeit jnp.dot(xjnp, xjnp.T).block_until_ready()
# 2.05 s ± 74.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.dot(xnp, xnp.T)
# 1.73 s ± 383 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  

Похоже, что в Numpy оптимизирована точечная операция для 32-разрядных чисел с плавающей точкой.