uint8 numpy数组的内存效率绝对差

Question

我有两个大 np.uint8 ndarrays，a 和 b。我需要计算：c = np.sum(np.abs(a - b), axis=(-2,-1,))

因为它们是无符号的，我不能只减去它们。一种天真的解决方法是将它们转换为更大的数据类型：

c = np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)), axis=(-2,-1,))

总共使用数组内存的 4^* 倍。在一个理想的世界里，我想成为这样的人：

c = np.sum(np.abssub(a, b), axis=(-2,-1,))

这将使用与数组相同的内存量。遗憾的是，我在 numpy 的文档中找不到这样的函数。现在我正在做以下事情：

diff = np.empty_like(a)

mask = a > b
diff[mask] = (a - b)[mask]
# b shape is different but broadcasts to a
# That is why I use mask after substracting

mask = np.logical_not(mask, out=mask)
diff[mask] = (b - a)[mask]

c = np.sum(np.abs(diff, out=diff), axis=(-2,-1,))

使用的是只是 2.5^**倍的内存量

有更好的方法吗？

*  4   times = bytes(a) + bytes(b) + bytes(a.astype(np.int16)) + bytes(b.astype(np.int16)) + bytes(a.astype(np.int16) - b.astype(np.int16))
               --------- 1 --------   ----------- 2 ----------  ----------- 3 -----------   --------------------- 4 ---------------------

** 2.5 times = bytes(a) + bytes(b) + bytes(diff) + bytes(mask) + bytes(a - b | b - a)
              --------- 1 --------   ------------ 2 ----------   ------- 2.5 -------

Answer 1

numexpr module 提供了一个非常简单但 memory-efficient 环境并且可以在这里使用。它会在执行算术运算时自动处理溢出问题。让我们看一个示例案例，看看如何解决我们的问题 -

In [63]: a = np.array([3,252,89],dtype=np.uint8)
    ...: b = np.array([10,255,19],dtype=np.uint8)

In [64]: import numexpr as ne

In [65]: ne.evaluate('abs(a-b)')
Out[65]: array([ 7.,  3., 70.])

因此，要获得所需的输出 -

In [66]: int(ne.evaluate('sum(abs(a-b))'))
Out[66]: 80

与 up-casted NumPy 版本比较 -

In [67]: np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))
Out[67]: 80

内存效率

现在，让我们扩展到一个非常大的数组，并检查问题的症结所在，即memory-efficiency。我们将使用 memory_profiler 模块进行相同的测试。

Python 脚本，其中 NumPy 和 numexpr 版本列为 numpy_numexpr_memeff.py -

import numpy as np
import numexpr as ne
from memory_profiler import profile

np.random.seed(0)
a = np.random.randint(0,256,(1000000))
b = np.random.randint(0,256,(1000000))

@profile(precision=10)
def numpy1():    
    return np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))

@profile(precision=10)
def numexpr():
    return int(ne.evaluate('sum(abs(a-b))'))

if __name__ == '__main__':
    numpy1()

if __name__ == '__main__':
    numexpr()

脚本 command-line 运行的结果 -

$ python -m memory_profiler numpy_numexpr_memeff.py 
Filename: numpy_numexpr_memeff.py

Line #    Mem usage    Increment   Line Contents
================================================
     9  63.0468750000 MiB   0.0000000000 MiB   @profile(precision=10)
    10                             def numpy1():    
    11  65.3437500000 MiB   2.2968750000 MiB       return np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))


Filename: numpy_numexpr_memeff.py

Line #    Mem usage    Increment   Line Contents
================================================
    13  65.3437500000 MiB   0.0000000000 MiB   @profile(precision=10)
    14                             def numexpr():
    15  65.5859375000 MiB   0.2421875000 MiB       return int(ne.evaluate('sum(abs(a-b))'))

因此，与 NumPy 版本相比，numexpr 版本似乎占用了 1/10 的内存。

性能

计时 -

In [68]: np.random.seed(0)
    ...: a = np.random.randint(0,256,(1000000))
    ...: b = np.random.randint(0,256,(1000000))

In [71]: %timeit np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))
3.99 ms ± 88.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [72]: %timeit int(ne.evaluate('sum(abs(a-b))'))
4.71 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

因此，就性能而言，numexpr 版本接近，但不如 NumPy。

另一个人可能会利用这样一个事实，即如果我们输入一个 up-scaled 一个，另一个将在执行算术运算时自动 up-scaled 。所以，我们可以简单地做 -

np.sum(np.abs(a.astype(np.int16) - b))

Python 用于测试 memory-efficiency 的脚本，如 numpys_memeff.py -

import numpy as np
from memory_profiler import profile

np.random.seed(0)
a = np.random.randint(0,256,(1000000))
b = np.random.randint(0,256,(1000000))

@profile(precision=10)
def numpy1():    
    return np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))

@profile(precision=10)
def numpy2():    
    return np.sum(np.abs(a.astype(np.int16) - b))

if __name__ == '__main__':
    numpy1()

if __name__ == '__main__':
    numpy2()

结果 -

$ python -m memory_profiler numpys_memeff.py 
Filename: numpys_memeff.py

Line #    Mem usage    Increment   Line Contents
================================================
     8  56.6015625000 MiB   0.0000000000 MiB   @profile(precision=10)
     9                             def numpy1():    
    10  59.1210937500 MiB   2.5195312500 MiB       return np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))


Filename: numpys_memeff.py

Line #    Mem usage    Increment   Line Contents
================================================
    12  59.1210937500 MiB   0.0000000000 MiB   @profile(precision=10)
    13                             def numpy2():    
    14  59.3632812500 MiB   0.2421875000 MiB       return np.sum(np.abs(a.astype(np.int16) - b))

在性能上，似乎也稍微好一些 -

In [68]: np.random.seed(0)
    ...: a = np.random.randint(0,256,(1000000))
    ...: b = np.random.randint(0,256,(1000000))

In [71]: %timeit np.sum(np.abs(a.astype(np.int16) - b.astype(np.int16)))
3.99 ms ± 88.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [73]: %timeit np.sum(np.abs(a.astype(np.int16) - b))
3.84 ms ± 29.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Answer 2

您可以使用

为自己节省一些中间数组

# sizeof(a)
diff = a - b
# sizeof(a)
mask = b > a
np.negative(diff, where=mask, out=diff)

c = np.sum(diff, axis=(-2,-1,))

或另一种拼写方式：

def abssub(a, b):
    diff = a - b
    mask = b > a
    return np.negative(diff, where=mask, out=diff)

c = np.sum(abssub(a, b), axis=(-2,-1,))

uint8 numpy数组的内存效率绝对差

Memory-efficient absolute difference of uint8 numpy arrays

python

arrays

numpy

integer-overflow