Question

背景

我有一个2D numpy数组，它表示大量的网格坐标矢量，并且每个坐标矢量都需要转换为字节字符串，以便可以将其转换为python集。

此字节字符串转换过程是代码运行时的真正瓶颈，因此，我正在寻找加快速度的方法。

示例代码

from numpy import int16
from numpy.random import randint
# make an array of coordinate vectors full of random ints
A = randint(-100,100,size = (10000,5), dtype=int16)
# pull each vector out of the array using iteration and convert to byte string
A = [v.tobytes() for v in A]
# build a set using the byte strings
S = set(A)

计时测试

使用timeit测试我们得到的当前代码

setup = 'from numpy import int16; from numpy.random import randint; A = randint(-100,100,size = (10000,5), dtype=int16)'
code = 'S = set([v.tobytes() for v in A])'
t = timeit(code, setup = setup, number=500)
print(t)
>>> 1.136594653999964

转换后构建集合的时间少于总计算时间的15％：

setup = 'from numpy import int16; from numpy.random import randint; A = randint(-100,100,size = (10000,5), dtype=int16); A = [v.tobytes() for v in A]'
code = 'S = set(A)'
t = timeit(code, setup = setup, number=500)
print(t)
>>> 0.15499859599980482

还值得注意的是，将整数的大小加倍至32位只会对运行时间产生很小的影响：

setup = 'from numpy import int32; from numpy.random import randint; A = randint(-100,100,size = (10000,5), dtype=int32)'
code = 'S = set([v.tobytes() for v in A])'
t = timeit(code, setup = setup, number=500)
print(t)
>>> 1.1422132620000411

这使我怀疑这里的大多数时间都被迭代或对tostring()的函数调用的开销所吞噬了。

如果是这种情况，我想知道是否有矢量化的方法可以避免迭代？

谢谢！

Answer 1

这是使用np.frombuffer的矢量化方法-

# a : Input array of coordinates with int16 dtype
S = set(np.frombuffer(a,dtype='S'+str(a.shape[1]*2)))

给定样本数据集的时间-

In [83]: np.random.seed(0)
    ...: a = randint(-100,100,size = (10000,5), dtype=int16)

In [128]: %timeit set([v.tobytes() for v in a])
2.71 ms ± 99.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [129]: %timeit set(np.frombuffer(a,dtype='S'+str(a.shape[1]*2)))
933 µs ± 4.16 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [130]: out1 = set([v.tobytes() for v in a])

In [131]: out2 = set(np.frombuffer(a,dtype='S'+str(a.shape[1]*2)))

In [132]: (np.sort(list(out1))==np.sort(list(out2))).all()
Out[132]: True

向量化2D numpy数组中的列到字节字符串的转换

1 个答案: