JAX 的运行速度是否比 NumPy 慢？答案

【问题标题】：Does JAX run slower than NumPy?JAX 的运行速度是否比 NumPy 慢？
【发布时间】：2022-01-17 23:05:19
【问题描述】：

我最近开始学习 JAX。我用 NumPy 写了一个简短的 sn-p 并用 JAX 写了它的等价物。我原以为 JAX 会更快，但是当我分析代码时，NumPy 代码比 JAX 代码快得多。我想知道这是否普遍正确，或者我的实现是否存在问题。

NumPy 代码：

import numpy as np
from time import time as tm


def gp(x):
    return np.maximum(np.zeros(x.shape), x)


# -- inputs
n_q1 = 25
n_q2 = 5

x = np.random.rand(1, n_q1)  # todo: changes w/ time
y = np.random.rand(1, n_q2)  # todo: changes w/ time

# -- activations
n_p1 = 3
n_p2 = 2

v_p1 = np.random.rand(1, n_p1)
v_p2 = np.random.rand(1, n_p2)

a_p1 = 0.5
a_p2 = 0.5

# -- weights
W_q1p1 = np.random.rand(n_q1, n_p1)
W_p2q2 = np.random.rand(n_p2, n_q2)
W_p1p1 = np.random.rand(n_p1, n_p1)
W_p1p2 = np.random.rand(n_p1, n_p2)
W_p2p1 = np.random.rand(n_p2, n_p1)

# -- computation
t1=tm()

for t in range(2000):

    z_p1 = np.matmul(v_p1, W_p1p1) + np.matmul(v_p2, W_p2p1) + np.matmul(x, W_q1p1)
    v_p1_new = a_p1 * v_p1 + (1 - a_p1) * gp(z_p1)

    z_p2 = np.matmul(v_p1, W_p1p2)
    v_p2_new = a_p2 * v_p2 + (1 - a_p2) * gp(z_p2)

    v_p1, v_p2 = v_p1_new, v_p2_new

print(tm()-t1)

这会产生：0.02118515968322754

JAX 代码：

from jax import random, nn, numpy as jnp

from time import time as tm


def gp(x):
    return nn.relu(x)


# -- inputs
n_q1 = 25
n_q2 = 5
key = random.PRNGKey(0)

x = random.normal(key, (1, n_q1))
y = random.normal(key, (1, n_q2))  # todo: check if I need to advance "key" manually

# -- activations
n_p1 = 3
n_p2 = 2

v_p1 = random.normal(key, (1, n_p1))
v_p2 = random.normal(key, (1, n_p2))

a_p1 = 0.5
a_p2 = 0.5

# -- weights
W_q1p1 = random.normal(key, (n_q1, n_p1))
W_p2q2 = random.normal(key, (n_p2, n_q2))

W_p1p1 = random.normal(key, (n_p1, n_p1))
W_p1p2 = random.normal(key, (n_p1, n_p2))
W_p2p1 = random.normal(key, (n_p2, n_p1))

# -- computation
t1=tm()
for t in range(2000):

    z_p1 = jnp.matmul(v_p1, W_p1p1) + jnp.matmul(v_p2, W_p2p1) + jnp.matmul(x, W_q1p1)
    v_p1_new = a_p1 * v_p1 + (1 - a_p1) * gp(z_p1)

    z_p2 = jnp.matmul(v_p1, W_p1p2)
    v_p2_new = a_p2 * v_p2 + (1 - a_p2) * gp(z_p2)

    v_p1, v_p2 = v_p1_new, v_p2_new
    
print(tm()-t1)

这会产生：2.5548229217529297

【问题讨论】：

标签： python numpy jax

【解决方案1】：

JAX 文档在其常见问题解答中有一个有用的部分：https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy

TL;DR：这很复杂。对于 CPU 上的单个矩阵运算，JAX 通常比 NumPy 慢，但 JAX 中 JIT 编译的操作序列通常比 NumPy 快，一旦转向 GPU/TPU，JAX 通常会比 NumPy 快得多。

【讨论】：