Индексирование в тензорном потоке медленнее, чем собирать

Question

Я пытаюсь индексировать в тензор, чтобы получить срез или отдельный элемент из 1d тензоров. Я обнаружил, что при использовании numpy способ индексации [:] а также slice vs tf.gather (почти 30-40%).

Также я наблюдаю, что tf.gather имеет значительные накладные расходы при использовании в скалярах (зацикливание на не сложенном тензоре) в отличие от тензора. Это известная проблема?

пример кода (неэффективно):

for node_idxs in graph.nodes():
    node_indice_list = tf.unstack(node_idxs)
    result = []
    for nodeid in node_indices_list:
        x = tf.gather(..., nodeid)
        y = tf.gather(..., nodeid)
        result.append(tf.mul(x,y))
return tf.stack(result)

в отличие от примера кода (эффективный):

for node_idxs in graph.nodes():
    x = tf.gather(..., node_idxs)
    y = tf.gather(..., node_idxs)
return tf.mul(x, y)

Я понимаю, что первая неэффективная реализация проделывает большую работу по разборке стеков, стеков, циклов и большему количеству операций по сбору, но я не ожидал 100-кратного замедления, когда порядок узлов, на которых я работаю, составляет несколько сотен узлов (это не стекирование и накладные расходы на сборку на одном скаляре, который медленен, в первом случае у меня есть намного больше операций по сбору, каждая из которых работает на одном элементе, в отличие от тензора смещений). Есть ли более быстрый способ индексации, я пробовал NumPy и Slice, который оказался медленнее, чем собирать.

Answer 1

Во-первых, код на самом деле не сравнивает сбор и индексирование Numpy — он сравнивает векторизованное индексирование (tf.gather) с циклическим индексированием (цикл Python «for»). Неудивительно, что зацикливание происходит медленно.

Обратите внимание, что индексация в стиле Numpy tensor[idxs]в любом случае ограничено в Tensorflow:

Только целые числа, срезы ( :), многоточие ( ...), tf.newaxis ( None) и скалярные тензоры tf.int32/tf.int64 являются действительными индексами

Так что используйте tf.gatherдля общих приложений.