Как эффективно извлечь скользящее окно в Tensorflow? - PullRequest
Новая
       10

Как эффективно извлечь скользящее окно в Tensorflow?

0 голосов
/ 25 октября 2019

Допустим, у меня есть вектор x с размерностью [n, n, 2n-1, 2n-1], и я хочу получить вектор y с размерностью [n, n, n, n] по следующим правилам: 

for i in range(n):
  for j in range(n):
    y[i, j] = x[n-i-1:2*n-i-1, n-j-1:2*n-j-1]

Например, скажем, x равен 

[[[[ 0  1  2]
   [ 3  4  5]
   [ 6  7  8]]

  [[ 9 10 11]
   [12 13 14]
   [15 16 17]]]


 [[[18 19 20]
   [21 22 23]
   [24 25 26]]

  [[27 28 29]
   [30 31 32]
   [33 34 35]]]]

Я хочу, чтобы вывод был 

[[[[ 4  5]
   [ 7  8]]

  [[12 13]
   [15 16]]]


 [[[19 20]
   [22 23]]

  [[27 28]
   [30 31]]]]

В np я могу сделать это эффективно, используя расширенную индексацию: 

n = 2
x = np.arange(n * n * (2*n-1) * (2*n-1)).reshape(n, n, 2*n-1, 2*n-1)
in_1 = np.arange(n)[:, np.newaxis, np.newaxis, np.newaxis]
in_2 = np.arange(n)[np.newaxis, :, np.newaxis, np.newaxis]
out_1 = (np.arange(n)[np.newaxis, np.newaxis, :, np.newaxis] + 
  (np.arange(n+1)[n:0:-1]-1)[:, np.newaxis, np.newaxis, np.newaxis])
out_2 = (np.arange(n)[np.newaxis, np.newaxis, np.newaxis, :] + 
  (np.arange(n+1)[n:0:-1]-1)[np.newaxis, :,np.newaxis, np.newaxis])
print(x[in_1, in_2, out_1, out_2])

Но в Tensorflow я не могу прийтис любым эффективным способом реализации этого. Я попытался использовать tf.gather_nd (используя индексацию, аналогичную описанной выше), а также tf.scan (в основном цикл for в качестве первого фрагмента кода), но оба они очень медленные. tf.image.extract_patches также не работает, как для каждого входного канала.

Для случая 1d (то есть вектор [n, 2n-1] для вектора [n, n], использующего аналогичное скольжение), tf.linalg.diag, за которым следуетТранспонирование может эффективно достичь вышеуказанного, но его нельзя обобщить на 2d случай.

...