[模型部署] ONNX模型转TRT模型部分要点

本篇讲“ONNX模型转TRT模型”部分要点。
以下皆为TRT模型的支持情况。
模型存为ONNX格式后，已经丢失了很多信息，与原python代码极难对应。因为在“ONNX转TRT”时，转换出错，更难映射回python代码。
解决此类问题的关键为：

转onnx时要打开verbose选项，输出每一行python的模型代码被转成了哪些ONNX算子。

    torch.onnx.export(
        model,
        (dummy_inputs),
        f=onnx_file_path,
        input_names=list(dummy_inputs.keys()),
        output_names=output_names,
        opset_version=...,
        verbose=True, 
    )
1
2
3
4
5
6
7
8
9

TRT 模型算子支持 （opset11）

where

where操作的每一项必须为bool

切片

切片的索引不能是bool，必须是int/long，且long的适用范围远多于int

Onehot

TRT不支持onehot操作

Nan

不支持"IsNan"，即进模型起不能有nan值

not操作

不支持tensor的not操作

bool的tensor不支持单个坐标切片

subA = A[:, 0] # A [B,N,D]
会出错误“Theis version of TensorRT does not support BOOL input type for the Gather operator”

方案：bool的tensor转long，切片后再转回bool

使用bool的tensor作mask时，tensor的维度必须与主tesnor一致

如 A[A_mask] = 0 # A_mask.dim() 会出现错误“Where_***. broadcast dimensions must be conformable”

方案:临时将A_mask扩展为A的尺寸。