python - 为什么 jax.numpy.dot() 在 CPU 上运行比 numpy.dot() 慢?
问题描述
我想使用 JAX 在 CPU 上加速我的 numpy 代码,然后在 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 是否也应该加速 CPU 上的 numpy 代码?
解决方案
Jax 和 Numpy 之间可能存在性能差异,但在原始帖子中,时间差异主要归结为数组创建中的错误。Jax 使用的数组具有 3000x3000 的形状,而 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
在测量定时性能时,应该收集多次运行,因为函数的性能是时间的分布而不是单个值。这可以通过 Python 标准库timeit.timeit
函数或%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 位浮点数的优化点操作。
推荐阅读
- c# - 将一个文本文件中的字符串行与另一个文本文件进行比较,并在发生不匹配时显示错误
- javascript - 如果一个函数不使用等待,它是否应该是异步的
- objective-c - 带有 NSFetchedResultsController 的 UICollectionViewController 坏了,当更新所有对象的属性时,如果控制器按此属性排序
- django - 如何将数据从 get_avatar()(Google 登录)发送到我的模板?
- python - 使用 python PIL 调整文件夹中的图像大小
- python - 编写返回列表中最短单词的函数
- sass - 为什么 `gulp-sass` 看不到我的 `mixins.scss` 文件?
- javascript - 未捕获的语法错误:意外的标记“?”
- python - spacy train 后的最终模型是否使用整个数据创建?
- python - Django:插入或更新数据库条目