张量流中的索引比收集慢答案

【问题标题】：indexing in tensorflow slower than gather张量流中的索引比收集慢
【发布时间】：2018-02-13 08:26:46
【问题描述】：

我正在尝试索引张量以从一维张量中获取切片或单个元素。我发现使用numpy 索引[:] 和slice vs tf.gather 的方式存在显着的性能差异（几乎30-40%）。

我还观察到 tf.gather 在用于标量（循环未堆叠张量）而不是 tensor 时具有显着的开销。这是一个已知问题吗？

示例代码（低效）：

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，结果比收集慢。

【问题讨论】：

标签： tensorflow tensorflow-serving tensorflow-gpu tensor tensorflow-xla

【解决方案1】：

首先，代码并没有真正比较收集与 Numpy 索引 - 它比较矢量化索引 (tf.gather) 与循环索引（Python“for”循环）。循环很慢也就不足为奇了。

请注意，类似 Numpy 的索引 tensor[idxs] 在 Tensorflow 中无论如何都受到限制：

仅整数、切片 (:)、省略号 (...)、tf.newaxis (None) 和标量 tf.int32/tf.int64 张量是有效的索引

所以对于一般应用，请使用tf.gather。

【讨论】：