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

Question

我有两个大的 np.uint8 ndarray，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 提供了一个非常简单但内存高效的环境，可以在这里使用。在执行算术运算时，它会自动处理溢出问题。让我们看一个示例案例，看看如何解决我们的问题-

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

与升级后的 NumPy 版本进行比较 -

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

内存效率

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

将 NumPy 和 numexpr 版本列为 numpy_numexpr_memeff.py 的 Python 脚本 -

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()  

脚本的命令行运行结果 -

$ 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 版本。

---

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

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

Python 脚本来测试这个的内存效率，如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,))