在大型代码库中,我使用np.broadcast_to
来广播数组(这里只使用简单示例):
In [1]: x = np.array([1,2,3])
In [2]: y = np.broadcast_to(x, (2,1,3))
In [3]: y.shape
Out[3]: (2, 1, 3)
在代码的其他地方,我使用的第三方函数可以在Numpy数组上以矢量化方式运行,但这不是ufunc。这些函数不了解广播,这意味着在像y
这样的数组上调用这样的函数是低效的。诸如Numpy的vectorize
之类的解决方案并不好,因为虽然他们理解广播,但他们会在数组元素上引入for
循环,这样效率非常低。
理想情况下,我希望能够做的就是拥有一个我们可以调用的函数,例如: unbroadcast
,返回一个具有最小形状的数组,如果需要,可以将其广播回完整大小。例如:
In [4]: z = unbroadcast(y)
In [5]: z.shape
Out[5]: (1, 1, 3)
然后,我可以在z
上运行第三方函数,然后将结果广播回y.shape
。
有没有办法实现依赖于Numpy公共API的unbroadcast
?如果没有,是否有任何黑客可以产生预期的结果?
答案 0 :(得分:4)
我有一个可能的解决方案,所以会在这里发布(但如果有人有更好的解决方案,请随时回复!)。一种解决方案是检查数组的strides
参数,它在广播维度上为0:
def unbroadcast(array):
slices = []
for i in range(array.ndim):
if array.strides[i] == 0:
slices.append(slice(0, 1))
else:
slices.append(slice(None))
return array[slices]
这给出了:
In [14]: unbroadcast(y).shape
Out[14]: (1, 1, 3)
答案 1 :(得分:3)
这可能相当于您自己的解决方案,只是内置一点。它使用as_strided
中的numpy.lib.stride_tricks
:
import numpy as np
from numpy.lib.stride_tricks import as_strided
x = np.arange(16).reshape(2,1,8,1) # shape (2,1,8,1)
y = np.broadcast_to(x,(2,3,8,5)) # shape (2,3,8,5) broadcast
def unbroadcast(arr):
#determine unbroadcast shape
newshape = np.where(np.array(arr.strides) == 0,1,arr.shape) # [2,1,8,1], thanks to @Divakar
return as_strided(arr,shape=newshape) # strides are automatically set here
z = unbroadcast(x)
np.all(z==x) # is True
请注意,在我的原始答案中,我没有定义函数,结果z
数组的(64,0,8,0)
为strides
,而输入为(64,64,8,8)
。在当前版本中,返回的z
数组与x
具有相同的步幅,我想传递并返回数组会强制创建副本。无论如何,我们总是可以在as_strided
中手动设置步幅以在所有情况下获得相同的数组,但在上述设置中似乎没有必要。