python - Python 中类熵公式 (sum(xlogx)) 的高效计算

Question

我正在寻找一种计算向量熵的有效方法，无需对它们进行归一化，同时忽略任何非正值。

由于向量不是概率向量，不应归一化，所以我不能使用 scipy's entropy功能。

到目前为止，我找不到一个 numpy 或 scipy 函数来获得它，因此我的替代方案涉及将计算分为 2 个步骤，这涉及中间数组并减慢运行时间。如果任何人都可以想到用于此计算的单个函数，那么它将是交叉的。

下面是一个 timeit 脚本，用于测量我尝试过的几个备选方案。我正在使用预分配数组来避免在运行时重复分配和释放。可以通过设置 func_code 的值来选择要运行的备选方案。我包括了其中一个答案提供的 nansum。 My MacBook Pro 2019 的测量值是:matmul:16.720187613xlogy: 17.296380516南苏姆:20.059866123000003

import timeit

import numpy as np
from scipy import special


def matmul(arg):
    a, log_a = arg
    log_a.fill(0)
    np.log2(a, where=a > 0, out=log_a)
    return (a[:, None, :] @ log_a[..., None]).ravel()


def xlogy(arg):
    a, log_a = arg
    a[a < 0] = 0
    return np.sum(special.xlogy(a, a), axis=1) * (1/np.log(2))


def nansum(arg):
    a, log_a = arg
    return np.nansum(a * np.log2(a, out=log_a), axis=1)


def setup():
    a = np.random.rand(20, 1000) - 0.1
    log = np.empty_like(a)
    return a, log


setup_code = """
from __main__ import matmul, xlogy, nansum, setup
data = setup() 
"""
func_code = "matmul(data)"

print(timeit.timeit(func_code, setup=setup_code, number=100000))

Answer 1

在我的机器上，对数的计算花费了 matmul 大约 80% 的时间，因此它绝对是瓶颈，优化其他函数将导致可以忽略不计的加速。

坏消息是默认实现 np.log 尚未在大多数平台上优化。实际上，默认情况下它未矢量化，recent x86 Intel processors supporting AVX-512 除外。 (即基本上是服务器上的 Skylake 处理器和 PC 上的 IceLake 处理器，但不是最近的 AlderLake)。这意味着一旦矢量化，计算速度就会显着加快。 AFAIK，闭源 SVML 库确实支持 AVX/AVX2 并且可以加速它(仅在 x86-64 处理器上)。 SMVL 受 Numexpr 和 Numba 支持，它们可以更快，因为假设您可以访问非自由 SVML，它是 HPC 机器上经常可用的英特尔工具的一部分(例如 MKL、OneAPI 等)。

如果您无权访问 SVML，则还有两个可能的选项:

• 实现您自己优化的 SIMD log2 函数，这是可能的，但很难，因为它需要很好地理解硬件 SIMD 单元，当然还需要编写 C 或 Cython 代码。该解决方案包括将 log2 函数计算为 n 度多项式近似值(它可以精确到 1 ULP，n 虽然是一个一般不需要)。朴素的近似值(例如 n=1)实现起来非常简单，但对于科学用途来说往往过于不准确。
• 通常使用 Numba/Cython 实现多线程日志计算。这是一个绝望的解决方案，因为如果输入数据不够大，多线程会降低速度。

---

这是一个多线程 Numba 解决方案的例子:

import numba as nb

@nb.njit('(UniTuple(f8[:,::1],2),)', parallel=True)
def matmul(arg):
    a, log_a = arg
    result = np.empty(a.shape[0])
    for i in nb.prange(a.shape[0]):
        s = 0.0
        for j in range(a.shape[1]):
            if a[i, j] > 0:
                s += a[i, j] * np.log2(a[i, j])
        result[i] = s
    return result

这在我的 6 核 PC 上快了大约 4.3 倍(200 us VS 46.4 us)。但是，如果您在具有许多内核的服务器上运行如此小的数据集，您应该小心，因为它在某些平台上实际上可能会更慢。