1

zmap 函数的 scipy stats 模块有什么替代方法吗？我目前正在使用它来获取两个非常大的数组的 zmap 分数，这需要相当长的时间。

是否有任何库或替代品可以提高其性能？或者甚至是另一个获得 zmap 函数的作用？

您的想法和意见将不胜感激！

这是我下面的最小可重现代码：

from scipy import stats
import numpy as np

FeatureData = np.random.rand(483, 1)
goodData = np.random.rand(4640, 483)
FeatureNorm= stats.zmap(FeatureData, goodData)

这就是 scipy stats.zmap 在幕后所做的：

def zmap(scores, compare, axis=0, ddof=0):
    scores, compare = map(np.asanyarray, [scores, compare])
    mns = compare.mean(axis=axis, keepdims=True)
    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
    return (scores - mns) / sstd

关于如何针对我的用例优化它的任何想法？我可以使用像 numba 或 JAX 这样的库来进一步提升它吗？

4

1 回答 1

1

幸运的是，zmap代码非常简单。然而，numpy 的开销来自它必须实例化中间数组的事实。如果您使用诸如numbaor中可用的数值编译器jax，它可以融合这些操作并以更少的开销进行计算。

不幸的是，numba 不支持 and 的可选参数mean，std所以让我们看一下 JAX。作为参考，这里是 scipy 和函数的原始 numpy 版本的基准，在 Google Colab CPU 运行时计算：

import numpy as np
from scipy import stats

FeatureData = np.random.rand(483, 1)
goodData = np.random.rand(4640, 483)

%timeit stats.zmap(FeatureData, goodData)
# 100 loops, best of 3: 13.9 ms per loop

def np_zmap(scores, compare, axis=0, ddof=0):
    scores, compare = map(np.asanyarray, [scores, compare])
    mns = compare.mean(axis=axis, keepdims=True)
    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
    return (scores - mns) / sstd

%timeit np_zmap(FeatureData, goodData)
# 100 loops, best of 3: 13.8 ms per loop

以下是在 JAX 中执行的等效代码，包括 Eager 模式和 JIT 编译：

import jax.numpy as jnp
from jax import jit

def jnp_zmap(scores, compare, axis=0, ddof=0):
    scores, compare = map(jnp.asarray, [scores, compare])
    mns = compare.mean(axis=axis, keepdims=True)
    sstd = compare.std(axis=axis, ddof=ddof, keepdims=True)
    return (scores - mns) / sstd

jit_jnp_zmap = jit(jnp_zmap)

FeatureData = jnp.array(FeatureData)
goodData = jnp.array(goodData)
%timeit jnp_zmap(FeatureData, goodData).block_until_ready()
# 100 loops, best of 3: 8.59 ms per loop

jit_jnp_zmap(FeatureData, goodData)  # trigger compilation
%timeit jit_jnp_zmap(FeatureData, goodData).block_until_ready()
# 100 loops, best of 3: 2.78 ms per loop

JIT 编译的版本比 scipy 或 numpy 代码快大约 5 倍。在 Colab T4 GPU 运行时，编译后的版本再增加 10 倍：

%timeit jit_jnp_zmap(FeatureData, goodData).block_until_ready()
1000 loops, best of 3: 286 µs per loop

如果这种操作是您分析中的瓶颈，那么像 JAX 这样的编译器可能是一个不错的选择。

于 2020-10-22T23:41:21.653 回答