pytorch转onnx后使用TensorRT部署转换踩坑记录
在pytorch训练好的模型转换为onnx模型然后转为TensorRT模型时报错:
onnx_parser/onnx2trt_utils.hpp:412 In function convert_axis:
[8] Assertion failed: axis >= 0 && axis < nbDims
builder/trt_builder.cpp:172]:Can not parse OnnX file: model/model/smartengine_model/detection/tinynet.onnx
查阅pytorch文档,才发现:
1. Pytorch模型中不能有条件支路
因为Pytorch模型转onnx模型的过程中,需要传给模型一个假的输入,然后做一次前向传播,再根据前向传播的结果来做模型的转换
2. Pytorch中的某些操作在TensorRT中不支持
x = x.view(x.size(0), -1)不支持,换成x = x.flatten(1)
torch.unsqueeze(x[:, 0], 1)不支持,需要去掉
动态size是不支持的,比如h = x.size()[2],改成具体数值
在每个分类模型都有一个平铺的过程,官方实现都是x=x.view(x.size(),-1)。你要做的是将这一句改成x=x.flatten(1), 导出来的onnx自然没问题。
参考:
https://ask.csdn.net/questions/1480596