Я хотел бы знать рекомендуемый способ ожидания завершения работы графического процессора в режиме TensorFlow Eager.
Операции, которые расположены на устройстве с графическим процессором, выполняются асинхронно (я не смог найти это в документации TensorFlow, но это соответствует поведению).Это важно, например, при синхронизации операций GPU с использованием time.time()*, поскольку нам необходимо убедиться, что операции выполнены, прежде чем регистрировать время окончания.
Единственный способ найти работу 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)
Поскольку копирование данных в ЦП сопряжено с некоторыми накладными расходами, было бы неплохо иметь способ обеспечитьGPU завершил выполнение без копирования.Есть ли такой способ?Возможно, что-то вроде JAX block_until_ready () ?
* Я понимаю, что использование time.time() может быть не лучшим способом для измерения времени работы графического процессора в режиме Eager.