TensorFlow Lite GPU 对 python 的支持

Posted 2023-02-16

【中文标题】TensorFlow Lite GPU 对 python 的支持

【英文标题】：Tensorflow Lite GPU support for python

【发布时间】：2019-10-04 15:27:54

【问题描述】：

有人知道 Tensorflow Lite 是否支持 Python 的 GPU 吗？我看过 android 和 ios 的指南，但我没有遇到任何关于 Python 的东西。如果安装tensorflow-gpu并导入tensorflow.lite.python.interpreter，会自动使用GPU吗？

【问题讨论】：

是的！它会自动处理 GPU 上的数据。但是，如果您想在任何嵌入式设备中使用 tensorflow lite，那么 tensorflow 会提供 TensorFlow Lite GPU 委托。您可以阅读this 了解更多信息。

@kruxx 但该指南似乎并未暗示支持 Python。这是否意味着 TFLite 不支持 Python 的 GPU？

如果您将tf.device 中的设备从 CPU 更改为 GPU，您也可以使用相同的型号在 GPU 上运行。

@kruxx 我试过了，但我没有得到任何 GPU 活动。是通过 GPU 委托让 GPU 进行处理的唯一方法吗？如果是这样，设置 tf.device 是不够的。

@JohnM 关于这个话题有什么消息吗？你找到解决办法了吗？

【参考方案1】：

您可以强制在 GPU 上进行计算：

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

希望这会有所帮助。

【讨论】：

但我正在尝试在 GPU 上运行 tensorflow.lite.python.interpreter。如果我将它包裹在tf.device 下，我不会得到 GPU 活动。

此解决方案适用于 Tensorflow 本身，对于 TFLite 将不起作用。

【参考方案2】：

根据this线程，不是。

【参考方案3】：

一种解决方案是将 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”