TensorFlow:GPU上int32和int64索引的聚集性能差异

我运行了一个小测试来评估几个TF操作的性能，发现了一个有趣的现象。在多台机器上，使用

int64

索引与

int32

索引相比，我始终可以获得更好的操作性能。例如，如下所示，对于

gather

，如果

index\u dtype

是

np.int32

，则操作速度比

np.int64

慢2倍

我的小测试由一个顺序运行的

聚集

操作图组成（元组确保依赖于前一个操作的输出）：

我多次运行图表并平均时间（忽略第一次运行）

---

由于，

int32

是GPU计算的本机int类型，因此我预期会出现相反的行为。性能差异的原因可能是什么？

一个假设——GPU聚集运算要求GPU索引位于主机内存中。当您输入32个索引时，它们被放置在GPU上，然后执行op需要将索引从GPU传输到CPU。dynamic_stitch也有类似的功能，本期文章展示了如何对这种情况进行分析——很好，可能就是这样。

indices = np.asarray(indices, dtype=index_dtype)
params = tf.placeholder(dtype=tf.float32)

op = tf.gather(params, indices)
for _ in range(num_ops):
    op = tf.gather(tf.tuple([params, op])[0], indices)