TensorRT解析难点
第⼆届TensorRT 加速 AI 推理⽐赛，利⽤ TensorRT 加速 ASR 模型 WeNet。比赛完成，简单总结下。
ONNX导入TensorRT通常会预见以下几个问题。
1. ONNX模型转化验证
Polygraphy工具可用于验证ONNX模型是否被TRT原生支持。若支持，则可用命令行直接转换模型；若不支持，则生成不支持的部分（ONNX格式），可通过Netron查看具体为哪些节点。
|
1 2 3 4 5 6 7 8 9 10 |
# 确认 TensorRT 是否完全原生支持该 .onnx polygraphy inspect capability model.onnx # 输出信息：[I] Graph is fully supported by TensorRT; Will not generate subgraphs. # 生成一个含有 TensorRT 不原生支持的 .onnx，再次用 inspect capability 来确认 python getOnnxModel-NonZero.py polygraphy inspect capability model-NonZero.onnx > result-NonZero.txt # 产生目录 .results，包含网络分析信息和支持的子图(.onnx)、不支持的子图(.onnx) |
2. TRT不支持的算子
对于TRT不能解析的算子，一般可以通过功能相同的其他算子或者算子组合来替代。常用工具为Onnx-Graphsurgeon。该工具支持修改ONNX模型中的节点，如删除、替换、增加等。例如Wenet模型中encoder部分，某个Slice节点不能接收布尔类型输入输出，解决方法就是在该节点前后加入Cast节点，将布尔张量转为Int32类型。
3. 精度切换与混合精度
要实现更快的推理速度常见方式是将精度降为Fp16甚至int8。但往往降低精度，会导致最终输出误差超过可接受范围。此时可通过polygraphy逐层对比ONNX与TRT误差，找出误差较大的节点，在TensorRT Api中将该节点调回Fp32。
|
1 2 3 4 5 6 7 |
for i in range(network.num_layers): layer = network.get_layer(i) print(f'Network Layer {i}: {layer.name}, {layer.type}, {layer.precision}, is_set: {layer.precision_is_set}') if layer.name == 'Conv_37': layer.precision = trt.float32 config.flags = 1 << int(trt.BuilderFlag.STRICT_TYPES) |
当然，这是个折中的方法，终极解决方式还是用Plugin进行优化。
4. C++ Plugin性能优化
常见Plugin使用场景为：模型中频繁出现的算子模块，可以替换为性能更优的Plugin。例如Wenet中LayerNorm在Onnx中为多个算子的组合，用Onnx-Graphsurgeon将这些节点合并为一个LayerNorm节点。同时，书写Plugin，通过Cuda核函数实现LayerNorm的计算逻辑。在Onnx转TRT的过程中，TRT会自动寻找与LayerNorm的同名Plugin进行解析。
