数组的均值

Question

我有一个这样的结构：

a = np.array([[np.array([1,2,3]), np.array([2])], [np.array([0,1,2,3,4]), np.array([0,4])]])

这是一个2x2的结构，每个元素的大小可以是任意的。

我想要计算每个元素的平均值，针对a：

np.mean(a)

ValueError: operands could not be broadcast together with shapes (3) (5) 

我试着调整axis，但没有成功。我想得到一个新的np.array，它的值是[[2, 2], [2, 2]]。

总的来说，我希望能对a运行任何向量化的函数，方式是一样的。

该怎么做呢？我需要快速的代码，所以请避免使用显式的循环。

我能做到的最好是：

f = np.mean
result = np.zeros(a.shape)
for i in np.ndindex(a.shape):
  result[i] = f(a[i])

Answer 1

根据你的评论，你的元素数量比较少，但每个元素都很大。这意味着外层循环的速度并不重要，而内层循环的速度才是关键。

我们来看看具体的数字。你的外层数组最多有4个维度，每个维度的大小最多为10。这意味着最多有10000个元素。而且这些元素“相当大”——我们保守地认为每个元素的大小是50。所以，你总共有510000次循环。如果你能提高10000次外层循环的速度，代码的性能提升也不会超过2%。实际上，这个2%的假设是基于除了循环本身没有其他工作要做，这显然不成立。

所以，你的关注点放错了。只有500000次内层循环才重要。如果你能把数组的数组替换成一个维度更高的单一数组，并且用numpy来处理，可能会更快，但为了提升不到1%的性能而让代码变得复杂且难以理解，这样做就不太聪明了。你可以直接使用vectorize或列表推导的方式，这些方法简单明了。

---

与此同时：

我可能应该尝试并行计算，为每个矩阵元素使用一个线程。

并行处理在这里是个好主意。但不要使用线程，也不要为每个元素都用一个线程。

假设你有一台8核的机器，使用8个硬件线程意味着你可以把事情做得快近8倍。但如果你使用10000个硬件线程，并不意味着你能快10000倍，甚至8倍——你可能会花很多时间在上下文切换、分配和释放线程栈等操作上，结果反而比顺序执行还要慢。所以，你应该创建一个包含8个硬件线程的工作池，把10000个任务放到队列中，让这个池来处理。

另外，10000个任务可能太细了。每个任务都有一点开销，如果任务太小，你就会浪费太多时间在这些开销上。一个明显的批处理方法是按轴进行——比如有1000个任务，每个任务处理一行10个元素，或者100个任务，每个任务处理一个包含100个元素的二维数组。可以尝试不同的批处理大小，比如0、1、2和3个轴，看看哪个性能最好。如果差别很大，你可能还想尝试稍微复杂一点的批处理，比如分成3D数组，形状为3x10x10和4x10x10。

最后，虽然Python的线程是真正的操作系统（因此也是硬件）线程，但全局解释器锁（GIL）限制了同一时间只能有一个线程在运行Python代码。numpy做了一些工作来绕过这个限制，但并不完美（尤其是当你在numpy调用之间有很多开销时）。最简单的解决办法是使用multiprocessing，这样你可以有一个包含8个独立进程的池，而不是在一个进程中有8个线程。这会显著增加开销——你要么需要复制子数组（可以隐式地作为任务函数的参数和返回值，或者显式地通过管道），要么把它们放在共享内存中。但如果你的任务足够大，通常这不是问题。

---

下面是如何并行化你现有代码的一个示例：

f = np.mean
result = np.zeros(a.shape)
future_map = {}
with futures.ProcessPoolExecutor(max_workers=8) as executor:
    for i in np.ndindex(a.shape):
        future_map[executor.submit(f, a[i])] = i
    for future in futures.as_completed(future_map):
        i = future_map[future]
        result[i] = future.result()

显然，你可以简化它（例如，用dict推导替换提交循环），但我想让你清楚需要改变的地方。

另外，我使用了futures（需要Python 3.2及以上版本；如果你用的是2.7，可以从PyPI安装回溯版本），因为这样代码会简单一些；如果你需要更多灵活性，可以使用更复杂的multiprocessing库。

最后，我没有进行任何批处理——每个任务都是一个单独的元素——所以开销可能会比较大。

但可以先从这个开始，尽量简化到你觉得舒服的程度，然后再转换成使用1、2和3个轴的批处理，如前面所述。

Answer 2

这是我能做到的最好效果：

a = np.array([[np.array([1,2,3]), np.array([2])], [np.array([0,1,2,3,4]), np.array([0,4])]])
np.array([ np.mean(b) for b in a.ravel()]).reshape(a.shape)

输出结果

array([[ 2.,  2.],
       [ 2.,  2.]])

可以推广到其他函数，像这样：

def ravel_func(func,arr):
    return np.asarray([ func(part) for part in arr.ravel()]).reshape(arr.shape)

速度测试，感谢 Jaime

In [82]: timeit vmean(a)
10000 loops, best of 3: 65 us per loop

In [83]: timeit ravel_func(np.mean,a)
10000 loops, best of 3: 67 us per loop

Answer 3

我想你想要的是 numpy.vectorize，它基本上就像内置的 map 函数，但专门用于处理多维数组（ndarrays）。

>>> a = np.array([[np.array([1,2,3]), np.array([2])], [np.array([0,1,2,3,4]), np.array([0,4])]])
>>> vmean = np.vectorize(np.mean)
>>> vmean(a)
array([[ 2.,  2.],
       [ 2.,  2.]])