python - 无法使用 Python 在 GPU (Jetson Nano) 上运行 tflite 模型

Question

我有一个量化的 tflite 模型，我想在 Nvidia Jetson Nano 上对其进行基准测试。我使用 tf.lite.Interpreter() 方法进行推理。该过程似乎不在 GPU 上运行，因为 CPU 和 GPU 上的推理时间是相同的。
有没有办法使用 Python 在 GPU 上运行 tflite 模型？
我试图通过设置 tf.device() 方法来强制使用 GPU，但仍然不起作用。官方文档有一些称为 GPU 加速的委托(delegate)，但我似乎找不到任何适用于 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'])

Answer 1

根据 link，TFLite 不支持 Nvidia GPU