python - 当简单索引不是时，为什么 numpy.take 对 numpy.split 的结果很慢？

Question

考虑以下代码

import timeit
import numpy as np

MyArray = np.empty((10000, 10000, 1))
print((MyArray.size, MyArray.shape, MyArray.dtype, np.isfortran(MyArray)))
print(timeit.timeit(lambda: MyArray[0], number=10000))
print(timeit.timeit(lambda: MyArray.take(0), number=10000))

MyTwoArrays = np.empty((10000, 10000, 2))
MyArray = np.split(MyTwoArrays, 2, axis=2)[0]
print((MyArray.size, MyArray.shape, MyArray.dtype, np.isfortran(MyArray)))
print(timeit.timeit(lambda: MyArray[0], number=10000))
print(timeit.timeit(lambda: MyArray.take(0), number=1))

及其在我系统上的输出:

(100000000, (10000, 10000, 1), dtype('float64'), False)
0.05690493136299911
0.06236779451013045
(100000000, (10000, 10000, 1), dtype('float64'), False)
0.0569617025453055
1.6303121549025763

MyArray 的两个版本具有相同的大小、形状、数据类型和数据顺序。尽管如此，与使用相同结果的简单索引或 numpy .take 使用“原生”numpy 数组。

这是为什么呢？我可以解决这个问题吗？

更新:

这似乎与 View 有关:MyArray = MyArray.copy() 解决了这个问题。尽管如此，我还是很想知道为什么 [0] 的运行速度一样快，而 numpy.take 的查看速度会变慢。

另一个更新:

我注意到减速取决于数组维度(不是数组维度的数量，而是数组元素的数量)。对于单个第 0 个元素，我实现了长达 8 秒的访问时间。我发现这是这个问题最令人惊讶的方面。无论 numpy.take 在内部做什么，我都看不出为什么当索引较大时这种额外的“间接级别”的计算速度应该较慢。

第三次更新:

根据 @hpaulj 的评论，MyArray.take(0) 和 MyArray[0] 不等价，这里是一个更正的代码示例。 (我凭借我的 MATLAB 直觉犯了那个错误，并一度停止验证我的最小示例。我不想替换原始示例，因为 hpaulj 的答案可能取决于它。)

import timeit
import numpy as np

for UseSplit in (True, False):
    if UseSplit:
        print("Using split")
        MyDoubleArray = np.random.rand(5000, 5000, 2)
        MyArray = np.split(MyDoubleArray, 2, axis=2)[0]
    else:
        print("Not using split")
        MyArray = np.random.rand(5000, 5000, 1)

    print((MyArray.size, MyArray.shape, MyArray.dtype, np.isfortran(MyArray)))


    NumpyTaking = MyArray.take(0)
    DirectIndexing = MyArray.item(0)
    assert (NumpyTaking == DirectIndexing)

    print("Take 1")
    print(timeit.timeit(lambda: MyArray.take(0), number=1))
    print("Index 1")
    print(timeit.timeit(lambda: MyArray.item(0), number=1))


    NumpyTaking = MyArray.take(0, axis=2)
    DirectIndexing = MyArray[:, :, 0]
    assert (NumpyTaking == DirectIndexing).all()

    print("Take many")
    print(timeit.timeit(lambda: MyArray.take(0, axis=2), number=1))
    print("Index many")
    print(timeit.timeit(lambda: MyArray[:, :, 0], number=1))

在我的(其他)系统上有这个输出:

Using split
(25000000, (5000, 5000, 1), dtype('float64'), False)
Take 1
0.2260607502818708
Index 1
2.1667747519799052e-05
Take many
0.44334302173084994
Index many
0.0005971935325195243
Not using split
(25000000, (5000, 5000, 1), dtype('float64'), False)
Take 1
2.851019410510247e-05
Index 1
2.0527339755549434e-05
Take many
0.13906132276656846
Index many
1.444516501325488e-05

Answer 1

这个答案有点长而且令人费解，但我认为关键是，MyArray[0] 是两种构造中的 View 。 MyArray.take 在第二种情况 (split) 中复制一份。 copy 需要更长的时间。

---

这两个 Action 不等价:

In [302]: MyArray = np.ones((10000, 10000, 1))
In [303]: MyArray[0].shape
Out[303]: (10000, 1)
In [304]: MyArray.take(0).shape
Out[304]: ()

take 和 axis=None(默认)，对数组进行分解。指定轴返回与 MyArray[0,:,:] 相同的东西:

In [305]: MyArray.take(0,axis=0).shape
Out[305]: (10000, 1)

使用 ipython timeit(和 numpy 1.14)

In [306]: timeit MyArray[0].shape
425 ns ± 7.03 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [307]: timeit MyArray.take(0).shape
1.25 µs ± 11.2 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [308]: timeit MyArray.take(0,axis=0).shape
10.6 µs ± 22.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

我有点惊讶 take 如此之慢，尽管我从来没有觉得它是一个速度工具。相反，它为以下情况提供了便利:

In [311]: MyArray.take(0, axis=1).shape
Out[311]: (10000, 1)
In [313]: MyArray[:,0,:].shape
Out[313]: (10000, 1)

在代码中使用数字而不是冒号指定轴更容易。

however, it can be easier to use if you need elements along a given axis.

当我通过拆分构造 MyArray 时，take 时间变得更糟

In [321]: timeit MyArray[0].shape
422 ns ± 5.54 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
In [322]: timeit MyArray.take(0).shape
713 ms ± 10.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [323]: timeit MyArray.take(0,axis=0).shape
705 ms ± 3.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

ravel 是这段额外时间的大部分时间。我认为 take 必须复制一份:

In [324]: timeit MyArray.ravel()
710 ms ± 19.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

构建步骤:np.ones((10000, 10000, 2)) 花费的时间更长，而且我担心我会出现内存错误。我使用 ones 而不是 empty 来确保数组在使用前已完全分配。

这表明内存管理问题使时序复杂化。

---

数据缓冲区指针告诉我数组是否是 View :

In [334]: MyArray.__array_interface__['data']
Out[334]: (139737581203472, False)
In [335]: MyArray2.__array_interface__['data']
Out[335]: (139737581203472, False)

MyArray2 就像您的 MyTwoArrays。所以拆分返回的是 View ，而不是副本。

但是 ravel 必须在 split case 中复制一份:

In [336]: MyArray.ravel().__array_interface__['data']
Out[336]: (139739981209616, False)
In [337]: MyArray2.ravel().__array_interface__['data']
Out[337]: (139737581203472, False)

查看索引与 take 的数据缓冲区:

In [343]: MyArray[0].__array_interface__['data']
Out[343]: (139737581203472, False)
In [344]: MyArray.take(0, axis=0).__array_interface__['data']
Out[344]: (34066048, False)
In [345]: MyArray.take(0).__array_interface__['data']
Out[345]: (33320032, False)

MyArray[0] 仍然是一个 View ，因此相对较快。

take MyArray 是一个副本，有和没有轴。

In [346]: timeit MyArray.copy()
701 ms ± 1.87 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

在第一种情况下，我可能应该返回并检查复制，但此 session 的内存负载正在拖累我的其余计算。