为什么这个函数在 JAX 和 numpy 中比较慢？答案

【问题标题】：Why is this function slower in JAX vs numpy?为什么这个函数在 JAX 和 numpy 中比较慢？
【发布时间】：2022-02-01 02:29:16
【问题描述】：

我有以下 numpy 函数，如下所示，我正在尝试使用 JAX 进行优化，但无论出于何种原因，它都会变慢。

有人能指出我可以做些什么来提高这里的性能吗？我怀疑这与 Cg_new 发生的列表理解有关，但将其分开并不会在 JAX 中产生任何进一步的性能提升。

import numpy as np 

def testFunction_numpy(C, Mi, C_new, Mi_new):
    Wg_new = np.zeros((len(Mi_new[:,0]), len(Mi[0])))
    Cg_new = np.zeros((1, len(Mi[0])))
    invertCsensor_new = np.linalg.inv(C_new)

    Wg_new = np.dot(invertCsensor_new, Mi_new)
    Cg_new = [np.dot(((-0.5*(Mi_new[:,m].conj().T))), (Wg_new[:,m])) for m in range(0, len(Mi[0]))] 

    return C_new, Mi_new, Wg_new, Cg_new

C = np.random.rand(483,483)
Mi = np.random.rand(483,8)
C_new = np.random.rand(198,198)
Mi_new = np.random.rand(198,8)

%timeit testFunction_numpy(C, Mi, C_new, Mi_new)
#1000 loops, best of 3: 1.73 ms per loop

这是 JAX 等价物：

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

def testFunction_JAX(C, Mi, C_new, Mi_new):
    Wg_new = jnp.zeros((len(Mi_new[:,0]), len(Mi[0])))
    Cg_new = jnp.zeros((1, len(Mi[0])))
    invertCsensor_new = jnp.linalg.inv(C_new)

    Wg_new = jnp.dot(invertCsensor_new, Mi_new)
    Cg_new = [jnp.dot(((-0.5*(Mi_new[:,m].conj().T))), (Wg_new[:,m])) for m in range(0, len(Mi[0]))] 

    return C_new, Mi_new, Wg_new, Cg_new

C = np.random.rand(483,483)
Mi = np.random.rand(483,8)
C_new = np.random.rand(198,198)
Mi_new = np.random.rand(198,8)

C = jnp.asarray(C)
Mi = jnp.asarray(Mi)
C_new = jnp.asarray(C_new)
Mi_new = jnp.asarray(Mi_new)

jitter = jax.jit(testFunction_JAX) 

%timeit jitter(C, Mi, C_new, Mi_new)
#1 loop, best of 3: 4.96 ms per loop

【问题讨论】：

标签： python performance numpy optimization jax

【解决方案1】：

有关 JAX 和 NumPy 之间基准比较的一般注意事项，请参阅https://jax.readthedocs.io/en/latest/faq.html#is-jax-faster-than-numpy

至于您的特定代码：当 JAX jit 编译遇到 Python 控制流（包括列表解析）时，它会有效地展平循环并分阶段执行完整的操作序列。这可能会导致 jit 编译时间变慢和代码不理想。幸运的是，您的函数中的列表理解很容易用原生 numpy 广播来表示。此外，您还可以进行另外两项改进：

在计算之前无需转发声明Wg_new 和Cg_new
在计算dot(inv(A), B) 时，使用np.linalg.solve 比显式计算逆运算更高效、更精确。

对 numpy 和 JAX 版本进行这三项改进会产生以下结果：

def testFunction_numpy_v2(C, Mi, C_new, Mi_new):
    Wg_new = np.linalg.solve(C_new, Mi_new)
    Cg_new = -0.5 * (Mi_new.conj() * Wg_new).sum(0)
    return C_new, Mi_new, Wg_new, Cg_new

@jax.jit
def testFunction_JAX_v2(C, Mi, C_new, Mi_new):
    Wg_new = jnp.linalg.solve(C_new, Mi_new)
    Cg_new = -0.5 * (Mi_new.conj() * Wg_new).sum(0)
    return C_new, Mi_new, Wg_new, Cg_new

%timeit testFunction_numpy_v2(C, Mi, C_new, Mi_new)
# 1000 loops, best of 3: 1.11 ms per loop
%timeit testFunction_JAX_v2(C_jax, Mi_jax, C_new_jax, Mi_new_jax)
# 1000 loops, best of 3: 1.35 ms per loop

由于改进了实现，这两个函数都比以前快了很多。但是，您会注意到，这里的 JAX 仍然比 numpy 慢。这在某种程度上是意料之中的，因为对于这种简单程度的函数，JAX 和 numpy 都有效地生成了在 CPU 架构上执行的相同的短系列 BLAS 和 LAPACK 调用。与 numpy 的参考实现相比，根本没有太大的改进空间，而且对于这么小的数组，JAX 的开销是显而易见的。

【讨论】：

哇，再次感谢您！优化的 numpy 版本确实派上用场了。
出于好奇，我对此进行了 numba 测试，它花费了 numpy 的 10 倍时间。哦，男孩。

【解决方案2】：

我使用perfplot 对一系列问题规模进行了测试。结果：jax 的速度稍微快了一点。 jax 在这里没有优于 numpy 的原因是它在 CPU 上运行（就像 NumPy 一样），而在这里，NumPy 已经非常优化。（它在引擎盖下使用 BLAS/LAPACK。）

重现情节的代码：

import jax.numpy as jnp
import jax
import numpy as np
import perfplot


def setup(n):
    C_new = np.random.rand(n, n)
    Mi_new = np.random.rand(n, 8)
    return C_new, Mi_new


def testFunction_numpy_v2(C_new, Mi_new):
    Wg_new = np.linalg.solve(C_new, Mi_new)
    Cg_new = -0.5 * (Mi_new.conj() * Wg_new).sum(0)
    return Wg_new, Cg_new


@jax.jit
def testFunction_JAX_v2(C_new, Mi_new):
    Wg_new = jnp.linalg.solve(C_new, Mi_new)
    Cg_new = -0.5 * (Mi_new.conj() * Wg_new).sum(0)
    return Wg_new, Cg_new


b = perfplot.bench(
    setup=setup,
    kernels=[testFunction_numpy_v2, testFunction_JAX_v2],
    n_range=[2 ** k for k in range(14)],
    equality_check=None
)
b.save("out.png")
b.show()

【讨论】：