python - 在 TensorFlow 的 Eager 模式下阻止 GPU 操作的最佳方法是什么？

Question

我想知道在 TensorFlow Eager 模式下等待 GPU 操作完成的推荐方法。

位于 GPU 设备上的操作似乎是异步执行的（我在 TensorFlow 文档中找不到这一点，但它与行为一致）。这很重要，例如，当使用time.time()* 计时 GPU 操作时，因为我们需要确保在记录结束时间之前完成操作。

我能找到确保 GPU 操作已执行的唯一方法是将（部分）输出数据显式复制到 CPU。

例如（假设所有操作都在 GPU 上进行）：

t0 = time.time()
result = f(input_tensor)  # carry out some operations on the input
_ = result[0].numpy()  # copies a single element of the output tensor to the CPU
t1 = time.time()
print("runtime =", t1 - t0)

由于将数据复制到 CPU 会产生一些开销，因此最好有办法确保 GPU 在不复制的情况下完成执行。有没有这样的方法？也许类似于 JAX 的block_until_ready()？

*我意识到time.time()在 Eager 模式下使用可能不是计时 GPU 操作的最佳方式。

标签： pythontensorflowgputiming