在numpy中有效填补张量

时间:2016-11-12 20:22:17

标签: python performance numpy vectorization

我有一个大小为X且类型为(n, m)的numpy数组np.uint8(因此它只包含[0, 255]中的值)。我还有从f[0, 255]的映射[0, 3]

我想创建一个形状为Y的数组(4, n, m),使其y_{k, i, j} = 1 if k == f(x_{i, j}),否则为0。现在,我这样做:

Y = np.zeros((4, n, m))
for i in range(256):
  Y[f(i), X == i] = 1

但这非常慢,而且我无法找到更有效的方法。有什么想法吗?

2 个答案:

答案 0 :(得分:1)

假设f可以一次性对所有迭代值进行操作,您可以使用broadcasting -

Yout = (f(X) == np.arange(4)[:,None,None]).astype(int)

运行时测试和验证 -

In [35]: def original_app(X,n,m):
    ...:     Y = np.zeros((4, n, m))
    ...:     for i in range(256):
    ...:         Y[f(i), X == i] = 1
    ...:     return Y
    ...: 

In [36]: # Setup Inputs
    ...: n,m = 2000,2000
    ...: X = np.random.randint(0,255,(n,m)).astype('uint8')
    ...: v = np.random.randint(4, size=(256,))
    ...: def f(x): 
    ...:     return v[x]
    ...: 

In [37]: Y = original_app(X,n,m)
    ...: Yout = (f(X) == np.arange(4)[:,None,None]).astype(int)
    ...: 

In [38]: np.allclose(Yout,Y)  # Verify
Out[38]: True

In [39]: %timeit original_app(X,n,m)
1 loops, best of 3: 3.77 s per loop

In [40]: %timeit (f(X) == np.arange(4)[:,None,None]).astype(int)
10 loops, best of 3: 74.5 ms per loop

答案 1 :(得分:1)

标量索引和布尔值的混合似乎会损害您的速度:

In [706]: %%timeit
     ...: Y=np.zeros((4,3,4))
     ...: for i in range(256):
     ...:   Y[f(i), X==i]+=1
     ...: 

100 loops, best of 3: 12.5 ms per loop

In [722]: %%timeit
     ...: Y=np.zeros((4,3,4))
     ...: for i in range(256):
     ...:     I,J=np.where(X==i)
     ...:     Y[f(i),I,J] = 1
     ...: 
100 loops, best of 3: 8.55 ms per loop

这是为了

X=np.arange(12,dtype=np.uint8).reshape(3,4)
def f(i):
    return i%4

在这种情况下,f(i)不是主要的时间消费者:

In [718]: timeit K=[f(i) for i in range(256)]
10000 loops, best of 3: 120 µs per loop

但获取X==i索引的速度很慢

In [720]: timeit K=[X==i for i in range(256)]
1000 loops, best of 3: 1.29 ms per loop
In [721]: timeit K=[np.where(X==i) for i in range(256)]
100 loops, best of 3: 2.73 ms per loop

我们需要重新考虑映射的X==i部分,而不是f(i)部分。

=====================

展平最后2个尺寸有助于;

In [780]: %%timeit 
     ...: X1=X.ravel()
     ...: Y=np.zeros((4,12))
     ...: for i in range(256):
     ...:     Y[f(i),X1==i]=1
     ...: Y.shape=(4,3,4)
     ...: 
100 loops, best of 3: 3.16 ms per loop