1 比较onnx模型和TensorRT模型推理结果的必要性

在进行模型部署的时候,如果我们使用TensorRT对模型加速,那么一般都会采用深度学习框架>onnx模型>TensorRT模型的技术路线,在onnx模型转换到TensorRT模型的过程中有可能某些算子会丢失精度,从而影响后续的模型推理结果,特别是对模型精度要求较高的问题(比如多分类问题),所以这个时候就需要有一个工具来比较onnx源模型与TensorRT模型的推理结果是否一致。

2 使用TensorRT的Polygraphy工具和onnx库比较推理结果

Polygraphy的Github仓库:https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy

在Polygraphy的Github仓库文档中为我们提供了一个比较二者结果的Python示例代码,

from polygraphy.backend.onnxrt import OnnxrtRunner, SessionFromOnnx
from polygraphy.backend.trt import TrtRunner, EngineFromNetwork, NetworkFromOnnxPath
from polygraphy.comparator import Comparator, DataLoader

model_path = "/path/to/model.onnx"

build_onnxrt_session = SessionFromOnnx(model_path)
build_engine = EngineFromNetwork(NetworkFromOnnxPath(model_path))

runners = [
    OnnxrtRunner(build_onnxrt_session),
    TrtRunner(build_engine),
]

data_loader = DataLoader(val_range=(0, 2))

run_results = Comparator.run(runners, data_loader=data_loader)
assert bool(Comparator.compare_accuracy(run_results))

主要的步骤如下:

  • 加载onnx模型
  • 从onnx模型中构建TensorRT模型
  • 将输入值限制在[0,2]的范围内
  • 分别使用onnx-runtime和TensorRT进行推理
  • 比较二者的推理结果是否一致