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

Posted 2023-02-16

【中文标题】无法使用 Python 在 GPU (Jetson Nano) 上运行 tflite 模型

【英文标题】：Cannot run tflite model on GPU (Jetson Nano) using Python

【发布时间】：2019-12-01 21:19:27

【问题描述】：

我有一个量化的 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'])

【参考方案1】：

根据link，TFLite 不支持 Nvidia GPU

【讨论】：

多解释一些会更好

TFLite 推理库目前不支持 Nvidia GPU。所以你不能在 Nvidia Jetson 上运行这样的 tflite 模型。需要将其转换成 tensorflow 或 tensorRT 模型才能在 Jetson 上执行

【参考方案2】：

根据this 最近的帖子，它似乎在 jetson nano 上可用。 但它看起来像自定义构建，请尝试使用它而不是 tensorflow lite。

如果您已经安装了它，如果该版本应该支持 GPU，可能会询问 nvidia developer。

或者您可以通过这种方式安装 nvidia 自定义 tensorflow。

Python 3.6+JetPack4.4

sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran
sudo apt-get install python3-pip
sudo pip3 install -U pip
sudo pip3 install -U pip testresources setuptools numpy==1.16.1 future==0.17.1 mock==3.0.5 h5py==2.9.0 keras_preprocessing==1.0.5 keras_applications==1.0.8 gast==0.2.2 futures protobuf pybind11
# TF-2.x
$ sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v44 tensorflow==2.2.0+nv20.8
# TF-1.15
$ sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v44 ‘tensorflow<2’

Python 3.6+JetPack4.3

$ sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev
$ sudo apt-get install python3-pip
$ sudo pip3 install -U pip
$ sudo pip3 install -U numpy grpcio absl-py py-cpuinfo psutil portpicker six mock requests gast h5py astor termcolor protobuf keras-applications keras-preprocessing wrapt google-pasta
# TF-2.x
$ sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v43 tensorflow==2.1.0+nv20.3
# TF-1.15
$ sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v43 tensorflow==1.15.2+nv20.3

【参考方案3】：

您的 Jetson Nano 是否支持 OpenCL？ 如果是这样，您可以将 OpenCL 委托与 TFLite 一起使用。 https://www.tensorflow.org/lite/guide/build_cmake#opencl_gpu_delegate