Używanie numpy.take do szybszego indeksowania fantazyjnego

Question

EDYCJA Utrzymałem bardziej skomplikowany problem, z którym stoję poniżej, ale moje problemy z np.take można podsumować lepiej w następujący sposób. Załóżmy, że masz tablicę img o kształcie (planes, rows) i inną macierz lut o kształcie (planes, 256) i chcesz użyć ich do utworzenia nowej tablicy out o kształcie (planes, rows), gdzie out[p,j] = lut[p, img[p, j]]. Można to osiągnąć z ozdobnego indeksowania następująco:

In [4]: %timeit lut[np.arange(planes).reshape(-1, 1), img] 
1000 loops, best of 3: 471 us per loop 

Ale jeśli zamiast ozdobnego indeksowania, należy użyć podjęcia i pętla pyton ciągu planes rzeczy można przyspieszyć ogromnie:

In [6]: %timeit for _ in (lut[j].take(img[j]) for j in xrange(planes)) : pass 
10000 loops, best of 3: 59 us per loop 

Can lut i img być w jakiś sposób uporządkowane, tak aby cała operacja odbyła się bez pętli python, ale przy użyciu numpy.take (lub alternatywnej metody) zamiast tradycyjnego wymyślnego indeksowania, aby utrzymać przewagę prędkości?

---

oryginalne pytanie Mam zestaw tabel przeglądowych (LUT), że chcą wykorzystać na obrazie. Tablica zawierająca LUT ma kształt (planes, 256, n), a obraz ma kształt (planes, rows, cols). Oba są z dtype = 'uint8', pasujące do osi LUT 256. Pomysł polega na uruchomieniu p -tej płaszczyzny obrazu poprzez każdą z LUTów n z p -tej płaszczyzny LUT.

Jeśli mój lut i img są następujące:

planes, rows, cols, n = 3, 4000, 4000, 4 
lut = np.random.randint(-2**31, 2**31 - 1, 
         size=(planes * 256 * n // 4,)).view('uint8') 
lut = lut.reshape(planes, 256, n) 
img = np.random.randint(-2**31, 2**31 - 1, 
        size=(planes * rows * cols // 4,)).view('uint8') 
img = img.reshape(planes, rows, cols) 

mogę osiągnąć to, co jestem po użyciu ozdobnego indeksowanie jak ten

out = lut[np.arange(planes).reshape(-1, 1, 1), img] 

co daje mi tablicę kształt (planes, rows, cols, n), gdzie out[i, :, :, j] przechowuje i -th płaszczyzny img przebieg przez j -th LUT z i -th płaszczyzny LUT ...

Wszystko jest dobre, z wyjątkiem tego:

In [2]: %timeit lut[np.arange(planes).reshape(-1, 1, 1), img] 
1 loops, best of 3: 5.65 s per loop 

który jest całkowicie nie do przyjęcia, zwłaszcza, że ​​mam wszystkie z poniższych nie tak ładne szukających alternatywy korzystających np.take niż działał znacznie szybciej:

1. Pojedyncza LUT na jednej płaszczyźnie działa o około x70 szybciej:

   In [2]: %timeit np.take(lut[0, :, 0], img[0]) 
   10 loops, best of 3: 78.5 ms per loop 
   

2. Pętla pyton uruchomiony przez wszystkich pożądanych kombinacji wykończenia niemal X6 szybciej:

   In [2]: %timeit for _ in (np.take(lut[j, :, k], img[j]) for j in xrange(planes) for k in xrange(n)) : pass 
   1 loops, best of 3: 947 ms per loop 
   

3. Nawet uruchomione wszystkie kombinacje samolotów w LUT i obrazu, a następnie odrzucając planes**2 - planes niechcianych te są szybsze niż fantazyjny indeksowania:

   In [2]: %timeit np.take(lut, img, axis=1)[np.arange(planes), np.arange(planes)] 
   1 loops, best of 3: 3.79 s per loop 
   

4. i najszybciej połączenie udało mi się wymyślić ma Iterowanie pętli pyton nad samolotami i kończy X13 szybciej:

   In [2]: %timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass 
   1 loops, best of 3: 434 ms per loop 
   

Pytanie brzmi, czy nie ma sposobu na zrobienie tego z np.take bez pętli Pythona? Idealnie, bez względu na to, czy potrzebne jest przekształcenie, czy zmiana rozmiaru, powinno to nastąpić na LUT, a nie na obrazie, ale jestem otwarty na to, co ludzie mogą wymyślić ...

Answer 1

Pięść z wszystkiego, co muszę powiedzieć, naprawdę polubiłem twoje pytanie. Bez przestawienie LUT lub IMG następujące rozwiązanie działa:

%timeit a=np.take(lut, img, axis=1) 
# 1 loops, best of 3: 1.93s per loop 

jednak od wyniku trzeba zapytania przekątnej: A [0,0] w [1,1] w [2,2]; aby dostać to, co chcesz. Próbowałem znaleźć sposób na indeksowanie tylko dla elementów przekątnych, ale nadal nie udało się.

Oto kilka sposobów, aby zmienić swoje LUT i IMG: następujące prace jeśli indeksy w IMG są od 0-255, w 1. płaszczyźnie, 256-511 dla 2 płaszczyzny i 512-767 dla 3 samolot, ale to by uniemożliwić korzystanie z 'uint8', który może być duży problem ...:

lut2 = lut.reshape(-1,4) 
%timeit np.take(lut2,img,axis=0) 
# 1 loops, best of 3: 716 ms per loop 
# or 
%timeit np.take(lut2, img.flatten(), axis=0).reshape(3,4000,4000,4) 
# 1 loops, best of 3: 709 ms per loop 

w moim komputerze rozwiązanie jest nadal najlepszym rozwiązaniem, i bardzo odpowiednie, ponieważ wystarczy ukośne oceny, tj. płaszczyzny1-płaszczyzny1, płaszczyzny2-płaszczyzny2 i płaszczyzny3-płaszczyzny3:

%timeit for _ in (np.take(lut[j], img[j], axis=0) for j in xrange(planes)) : pass 
# 1 loops, best of 3: 677 ms per loop 

Mam nadzieję, że to da ci pewien wgląd w lepsze rozwiązanie. Byłoby miło szukać więcej opcji z flatten() i podobnych metod, jak np.apply_over_axes() lub np.apply_along_axis(), które wydają się obiecujące.

użyłem tego poniższy kod, aby wygenerować dane:

import numpy as np 
num = 4000 
planes, rows, cols, n = 3, num, num, 4 
lut = np.random.randint(-2**31, 2**31-1,size=(planes*256*n//4,)).view('uint8') 
lut = lut.reshape(planes, 256, n) 
img = np.random.randint(-2**31, 2**31-1,size=(planes*rows*cols//4,)).view('uint8') 
img = img.reshape(planes, rows, cols)