本篇讲“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,
)
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()<A.dim()
会出现错误“Where_***. broadcast dimensions must be conformable”
方案:临时将
A_mask
扩展为A的尺寸。