TensorFlow Lite GPU 對 Python 的支持

Question

有人知道 Tensorflow Lite 是否支持 Python 的 GPU 嗎？ 我看過 Android 和 iOS 的指南，但我沒有遇到任何關於 Python 的內容。 如果安裝了tensorflow-gpu並導入了tensorflow.lite.python.interpreter ，會自動使用GPU嗎？

Answer 1

根據這個線程，它不是。

Answer 2

一種解決方案是將 tflite 轉換為 onnx 並使用 onnxruntime-gpu

使用https://github.com/onnx/tensorflow-onnx轉換為 onnx ：

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

然后pip install onnxruntime-gpu

並像這樣運行：

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

您可以通過以下方式獲取輸入和輸出名稱：

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

同樣，但將“get_inputs”替換為“get_outputs”

Answer 3

您可以強制計算在 GPU 上進行：

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

希望這可以幫助。