我想知道在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操作的最佳方法。