Почему тензор потока fftsthift (roll) работает так медленно? - PullRequest
Новая
       8

Почему тензор потока fftsthift (roll) работает так медленно?

0 голосов
/ 06 апреля 2020

Я профилировал свою сеть и понял, что операция Roll (в моем коде, используемом fftshift) отвечает за 25% времени устройства. Я продолжил профилировать эту операцию самостоятельно в блокноте, используя следующий код: 

import tensorflow as tf
from tensorflow.signal import fftshift, fft2d

test_tensor = tf.complex(tf.random.normal([1, 15, 640, 372]), tf.random.normal([1, 15, 640, 372]))

%%timeit
fftshift(test_tensor, axes=[2, 3])

1.79 ms ± 46.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%%timeit
fft2d(test_tensor)

442 µs ± 38.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Это было сделано на графическом процессоре Quadro P5000 с Python 3.6.8, tf 2.2.0rc2 и Ubuntu 16.04.

Результаты на ЦП немного менее странные с 1.32 ms ± 14.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) и 6.6 s ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) ( Интересно, что fftshift numpy по некоторым причинам медленнее в 20 раз.

Исходный код для операции поворота: https://github.com/tensorflow/tensorflow/blob/ddabed4285d27785213322d05dcbe0ebc392849d/tensorflow/core/kernels/roll_op.cc.

...