эффективно заполняет тензор в numpy

Question

У меня есть массив numbox X размера (n, m) и типа 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 

Но это супер медленно, и я не смог найти более эффективный способ сделать это. Есть идеи?

Answer 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 

Answer 2

Сочетание скалярной индексации и булево, как представляется, болеть вашей скоростью:

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