如何查询TensorFlow .pb冻结图、ONNX模型浮点精度并校验各转换环节精度

1. 各类模型的浮点精度查询方法 #

• 查询TensorFlow .pb冻结图精度
  你可以直接通过TensorFlow原生API加载冻结图，遍历常量节点的数据类型判断精度，参考代码如下：

  import tensorflow as tf
  from tensorflow.python.framework import tensor_util
  
  with tf.io.gfile.GFile("你的模型路径.pb", "rb") as f:
      graph_def = tf.compat.v1.GraphDef()
      graph_def.ParseFromString(f.read())
  
  for node in graph_def.node:
      if node.op == "Const":
          tensor = tensor_util.MakeNdarray(node.attr["value"].tensor)
          print(f"节点名: {node.name}, 数据类型: {tensor.dtype}")
  

  默认公开预训练的ResNet-50冻结图均为FP32精度，只有你手动做过混合精度训练、半精度导出操作时才会出现FP16类型节点。

• 查询ONNX模型精度
  借助ONNX官方库加载模型后遍历初始化张量即可查询精度，参考代码如下：

  import onnx
  
  model = onnx.load("你的模型路径.onnx")
  for init in model.graph.initializer:
      print(f"张量名: {init.name}, 数据类型: {onnx.helper.tensor_dtype_to_np_dtype(init.data_type)}")
  

  ONNX模型默认会完整保留源模型的精度类型，除非转换时手动指定了精度转换参数。

2. 各转换环节的浮点精度说明 #

• 阶段1：TensorFlow .pb → ONNX
  默认使用tf2onnx工具转换时，精度完全和源冻结图保持一致，工具不会主动修改浮点节点的类型。如果需要统一转为FP16精度，可以在转换命令中添加--dtype fp16参数强制转换所有浮点节点为FP16。
• 阶段2：ONNX → TensorRT Engine
  转换精度完全由你配置的TensorRT编译参数决定：
  • 默认配置下，TensorRT会完全沿用ONNX模型的源精度（通常为FP32）构建推理引擎
  • 开启FP16推理需要手动添加配置项：config.set_flag(trt.BuilderFlag.FP16)，此时TensorRT会自动将支持FP16运算的算子替换为半精度实现，不支持的算子保留原精度；如果需要强制所有算子都走FP16，可以额外添加config.set_flag(trt.BuilderFlag.STRICT_TYPES)配置
  • ResNet-50属于标准CV结构，转TensorRT FP16精度的推理误差极小，常规分类、检测场景下可以忽略，推理速度相比FP32提升1倍左右。

3. 注意事项 #

如果转换环节需要修改精度，建议每一步转换完成后都做相同输入的输出结果对齐验证，避免转换异常导致的精度损失问题。

内容的提问来源于stack exchange，提问作者Shruthi Manian