Python环境下无需全量TensorFlow依赖如何完成PB模型推理

你要实现和安卓端一致的、仅依赖推理runtime加载模型的效果，完全可以通过端侧专用推理框架实现，配合pyinstaller打包后体积可以从500MB压缩到几十MB级别，不需要引入全量TensorFlow包。

方案一：使用官方TFLite Runtime（优先推荐，适配性最好） #

这是TensorFlow官方推出的端侧推理专用runtime，剔除了所有训练、图构建、调试相关的冗余代码，是和安卓端TensorFlow Inference库定位完全一致的Python侧对应组件。

• 第一步：模型格式转换（仅需在开发环境执行一次，不需要打包进最终程序）
  你现有的.pb格式frozen graph需要先转成.tflite格式，转换时临时用到全量TensorFlow即可，转换完成后就可以彻底脱离全量TF依赖：

  # 转换阶段临时使用全量TensorFlow，最终打包程序不需要包含该部分代码
  import tensorflow as tf
  converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
      graph_def_file="your_model.pb",
      input_arrays=["你的输入张量名"],
      output_arrays=["你的输出张量名"],
      input_shapes={"你的输入张量名": [1, 输入高, 输入宽, 通道数]} # 替换为模型实际输入维度
  )
  # 若模型包含量化等自定义算子，可在此处开启对应算子支持
  tflite_model = converter.convert()
  with open("model.tflite", "wb") as f:
      f.write(tflite_model)
  

• 第二步：最终程序仅依赖tflite-runtime完成推理
  开发最终要打包的程序时，卸载全量TensorFlow，仅安装tflite-runtime包即可，该包全平台安装体积仅十几MB，无冗余依赖。推理代码示例：

  # 最终打包的代码中禁止出现import tensorflow相关语句
  import tflite_runtime.interpreter as tflite
  import numpy as np
  
  # 加载模型
  interpreter = tflite.Interpreter(model_path="model.tflite")
  interpreter.allocate_tensors()
  input_info = interpreter.get_input_details()
  output_info = interpreter.get_output_details()
  
  # 填充预处理后的输入数据
  input_data = np.random.rand(1, 224, 224, 3).astype(np.float32) # 替换为实际业务输入
  interpreter.set_tensor(input_info[0]['index'], input_data)
  # 执行推理
  interpreter.invoke()
  # 获取推理结果
  result = interpreter.get_tensor(output_info[0]['index'])
  

• 打包注意事项：
  用pyinstaller打包前先删除旧的build、dist目录和.spec缓存文件，打包时可以显式指定排除tensorflow依赖，避免缓存带入冗余文件，最终打包后的可执行文件体积通常在30-60MB区间。

方案二：使用ONNX Runtime（适合有跨框架部署需求的场景） #

如果你不想使用TFLite，也可以选择ONNX Runtime这个通用推理runtime，体积和tflite-runtime接近，同样无训练相关冗余。

• 开发环境下先通过tf2onnx工具将.pb模型转换为ONNX格式，该步骤同样仅需执行一次
• 最终程序仅安装onnxruntime包，直接加载转换后的ONNX模型完成推理即可，推理代码逻辑和TFLite类似，性能表现基本一致

不推荐为了直接加载原始.pb文件自行封装TensorFlow C API，该方案开发调试成本极高，且后续算子兼容维护成本远高于上述两个成熟方案。

内容的提问来源于stack exchange，提问作者Ali Ok