内存高效的uint8 numpy数组绝对差值

Question

我有两个很大的 np.uint8 数组，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 -------

- jabozzo

---

你是否愿意就地改变 a 和 b 的值以便重复使用它们的内存空间？ - Eric

np.copyto(diff, (b - a), where=mask) will use less memory than diff[mask] = (b - a)[mask] - Eric

Answer 1

numexpr模块提供了一个非常简单但内存使用高效的环境，可以在这里使用。它会自动处理算术运算时出现的溢出情况。让我们看一个示例并了解如何解决问题 -

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 模块来测试相同的内容。

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

在命令行下运行脚本后的结果 -

$ 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版本的内存占用只有十分之一。

性能

时间 -

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)

- Divakar

---

感谢您教我们如何使用 memory_profiler - 我迫不及待地想自己试试。 - Paul Panzer

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

- Eric