我有一个量化的 tflite 模型,我想在 Nvidia Jetson Nano 上进行基准测试。我使用 tf.lite.Interpreter() 方法进行推理。该过程似乎不在 GPU 上运行,因为 CPU 和 GPU 上的推理时间是相同的。
有没有办法使用 Python 在 GPU 上运行 tflite 模型?
我试图通过设置 tf.device() 方法来强制使用 GPU,但仍然不起作用。官方文档中有一些称为 GPU 加速的委托,但我似乎找不到任何适用于 Python 的东西。
with tf.device('/device:GPU:0'):
interpreter = tf.lite.Interpreter(model_path="model.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.uint8)
interpreter.set_tensor(input_details[0]['index'], input_data)
start_time = time.time()
interpreter.invoke()
elapsed_time = time.time() - start_time
print(elapsed_time)
output_data = interpreter.get_tensor(output_details[0]['index'])