深度学习模型调试器 polygraphy

对于深度学习模型的部署，往往涉及多种框架之间的转换，一般是训练框架 ( PyTorch、TensorFlow ... ) => 推理框架 ( TensorRT、OnnxRuntime ... )，每个框架都有不同的模型表示，所以这个过程中往往最关心的是转换之后不要掉精度。不过往往理想很丰满，现实不丰满，在出现掉精度的时候，总想有一个好用的工具能帮助我们定位是哪里出了问题，这个时候，polygraphy 就呼之欲出了。

1 polygraphy 介绍

polygraphy 是一个深度学习模型调试工具，包含 python API 和 命令行工具 ，关于 polygraphy 的相关介绍其实比较少，它有的一些功能如下：

使用多种后端运行推理计算，包括 TensorRT, onnxruntime, TensorFlow；
比较不同后端的逐层计算结果；
由模型恩建生成 TensorRT 引擎并序列化为.plan；
查看模型网络的逐层信息；
修改 Onnx 模型，如提取子图，计算图化简；
分析 Onnx 转 TensorRT 失败原因，将原计算图中可以 / 不可以转 TensorRT 的子图分割保存；
隔离 TensorRT 终端错误 tactic；

polygraphy 的安装方式比较简单，直接 pip 安装即可：

安装彩色日志：

pip install colored

pip install -i https://pypi.douban.com/simple nvidia-pyindex
pip install -i https://pypi.douban.com/simple polygraphy

2 polygraphy 使用示例

mmdeploy用法：

polygraphy surgeon sanitize end2end.onnx --fold-constants -o end2end_folded.onnx

示例代码：

这里介绍一个 polygraphy 使用的示例，对 onnxruntime 和 TensorRT 进行精度对比，流程差不多是这样的：

首先生成一个 .onnx 文件；
其次使用 polygraphy 生成一个 FP16 的 TRT 引擎，并对比使用 onnxruntime 和 TensorRT 的计算结果；
然后使用 polygraphy 生成一个 FP32 的 TRT 引擎，将网络中所有层都标记为输出，并对比使用 onnxruntime 和 TensorRT 的计算结果 (逐层结果对比)；

相关代码示意如下：

# 生成一个 .onnx 模型作为 polygraphy 的输入
# export model.onnx from pytorch
# or
# export model.onnx from tensorflow

# 使用上面生成的 model.onnx 构建 TensorRT 引擎，使用 FP16 精度同时在 TensorRT 和 onnxruntime 中运行
polygraphy run model.onnx \
    --onnxrt --trt \
    --workspace 100000000 \
    --save-engine=model_FP16.plan \
    --atol 1e-3 --rtol 1e-3 \
    --fp16 \
    --verbose \
    --trt-min-shapes 'x:0:[1,1,28,28]' \
    --trt-opt-shapes 'x:0:[4,1,28,28]' \
    --trt-max-shapes 'x:0:[16,1,28,28]' \
    --input-shapes 'x:0:[4,1,28,28]' \
    > result-run-FP16.txt

# 使用上面生成的 model.onnx 构建 TensorRT 引擎，使用 FP32 精度同时在 TensorRT 和 onnxruntime 中运行
# 输出所有层的计算结果作对比
polygraphy run model.onnx \
    --onnxrt --trt \
    --workspace 100000000 \
    --save-engine=model_FP32_MarkAll.plan \
    --atol 1e-3 --rtol 1e-3 \
    --verbose \
    --onnx-outputs mark all \
    --trt-outputs mark all \
    --trt-min-shapes 'x:0:[1,1,28,28]' \
    --trt-opt-shapes 'x:0:[4,1,28,28]' \
    --trt-max-shapes 'x:0:[16,1,28,28]' \
    --input-shapes 'x:0:[4,1,28,28]' \
    > result-run-FP32-MarkAll.txt

选取 FP16 推理对比结果进行展示，如下：

好了，以上分享了教你使用深度学习模型调试器 polygraphy。希望我的分享能对你的学习有一点帮助。