Cannot run tflite model on GPU (Jetson Nano) using Python

Question

I have a quantized tflite model that I'd like to benchmark for inference on a Nvidia Jetson Nano. I use tf.lite.Interpreter() method for inference. The process doesn't seem to run on the GPU as the inference times on both CPU and GPU are the same.

Is there any way to run a tflite model on GPU using Python?

I tried to force GPU usage by setting tf.device() method but still doesn't work. The official documentation has something called delegates for GPU acceleration but I can't seem to find anything for 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

根据此链接， TFLite 不支持 Nvidia GPU

Answer 2

It seems to be available on the jetson nano according to this recent thread. But it's look like a custom build, try it instead of tensorflow lite.

If you already installed it maybe ask for nvidia developper if the release is supposed to support GPU.

Or you can Install the nvidia custom tensorflow this way.

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

Answer 3

Does your Jetson Nano support OpenCL? If it does, you can use OpenCL delegate with TFLite. https://www.tensorflow.org/lite/guide/build_cmake#opencl_gpu_delegate

Answer 4

It is because you are some how using the full interpreter of tf. see here

interpreter = tf.lite.Interpreter(model_path="model.tflite").

What you can do is install

python3 -m pip install tflite-runtime

and use

import tflite_runtime.interpreter as tflite
interpreter = tflite.Interpreter(model_path=args.model_file)

and you should be able to run things fine.

I hope it helps!