python - 如何从 .pb 转换为 .tflite？

Question

我使用 Pytorch 创建了一个对象检测模型然后从 .pth 转换而来至 .onnx然后 .pb , 但现在我需要将其转换为 .tflite对于安卓应用程序!怎么做？这是我的第一次。

input_arrays  = [64, 3, 224, 224] 

output_arrays = ? 

用于二进制分类。

我是从 pytorch 完成的，但我发现要查看的所有内容都来自 keras 或 Tensorflow...

这是我用来从 .pb 转换过来的代码至 .tflie

converter = lite.TFLiteConverter.from_frozen_graph(
  model/model.pb , input_arrays, output arrays )
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

!tflite_convert \
   --output_file= model/model.tflite \
   --graph_def_file= model/model.pb \
   --input_arrays= input_arrays \
   -- output_arrays= output_arrays

我认为它与输入数组和输出数组有关，但不确定。是graph_def_file应该存储 model.pb ？

Answer 1

使用以下代码时无需指定输入输出数组:

import tensorflow as tf

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()
open("converted_model.tflite", "wb").write(tflite_model)

试试这个。