CODEWITHSUNDEEP
tensorflowmachine-learningtensorflow-lite
John M.
John M.

Reputation: 2712

Tensorflow Lite GPU support for python

Anyone know if Tensorflow Lite has GPU support for Python? I've seen guides for Android and iOS, but I haven't come across anything about Python. If tensorflow-gpu is installed and tensorflow.lite.python.interpreter is imported, will GPU be used automatically?

Upvotes: 12

Views: 4553

Answers (3)

Megan Hardy
Megan Hardy

Reputation: 397

one solution is to convert tflite to onnx and use onnxruntime-gpu

convert to onnx with https://github.com/onnx/tensorflow-onnx:

pip install tf2onnx
python3 -m tf2onnx.convert --opset 11 --tflite path/to/model.tflite  --output path/to/model.onnx

then pip install onnxruntime-gpu

and run like:

session = onnxruntime.InferenceSession(('/path/to/model.onnx'))
raw_output = self.detection_session.run(['output_name'], {'input_name': img})

you can get the input and output names by:

for i in range(len(session.get_inputs)):
    print(session.get_inputs()[i].name)

and the same but replace 'get_inputs' with 'get_outputs'

Upvotes: 1

RomanS
RomanS

Reputation: 923

According to this thread, it is not.

Upvotes: 3

Rajneesh Aggarwal
Rajneesh Aggarwal

Reputation: 397

You can force the computation to take place on a GPU:

import tensorflow as tf
with tf.device('/gpu:0'):
   for i in range(10):
         t = np.random.randint(len(x_test) )
         ...

Hope this helps.

Upvotes: -1

Related Questions