python3环境下pytorch模型转onnx部署tensorRT

转载请注明出处，谢谢！

环境：python3.5.2+pytorch1.2.0+tensorRT7.0.0.11

1、pytorch模型转onnx模型

pytorch自带onnx模型导出，几句话就可以实现很方便

model = model.cuda()
model = model.eval()
dummy_input = torch.randn(1, channel, height, width, device='cuda')
torch.onnx.export(model, dummy_input, onnx_path, verbose=True,input_names=input_names,output_names=output_names)

dummy_input是固定一个网络的输入，对于检测类的全卷积网络，模型的输入是不受限制的，所以这个输入尺寸在转onnx模型时是个可以随便设置的

torch.onnx.export是pytorch自带的onnx模型导出函数，onnx_path是onnx模型导出保存路径，input_names是网络输入名字，output_names是网络输出名字

onnx网络结构op划分比较细，网络结构的可视化可使用netron，使用方法可见官方github https://github.com/lutzroeder/netron

onnx可视化的网络结构看起来比较费劲，可使用一个网络简化模块来简化网络结构，简化后的网络就是我们正常熟悉的那些层了

onnx-simplifier的安装使用见官方github https://github.com/daquexian/onnx-simplifier

另外如果不使用onnx-simplifier，直接将onnx模型导入tensorRT时可能会报错

2、python环境下安装tensorRT

官网下载链接https://developer.nvidia.com/nvidia-tensorrt-download

下载合适自己电脑环境的RT版本，以TensorRT 7.0.0.11为例子

#解压下载好的压缩文件
tar xzvf TensorRT-7.0.0.11.Ubuntu-16.04.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz
#将lib绝对路径添加到环境变量中，使用vim ~/.bashrc
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/oyrq/TensorRT-7.0.0.11/lib
#安装TensorRT，这里匹配的是python3.5的环境
cd TensorRT-7.0.0.11/python
sudo pip install tensorrt-7.0.0.11-cp35-none-linux_x86_64.whl
#安装graphsurgeon
cd TensorRT-7.0.0.11/graphsurgeon
sudo pip install graphsurgeon-0.4.1-py2.py3-none-any.whl

import tensorrt进行测试就OK了

3、tensorRT导入onnx模型

相关代码可参考nvidia-tensorRT里面的例子，TensorRT-7.0.0.11/samples/python下面有很多种，主要参考yolov3_onnx/onnx_to_tensorrt.py就可以了

4、tensorRT导入onnx模型报错

1）、[TensorRT] ERROR: Parameter check failed at: ../builder/Layers.cpp::ConstantLayer::1585, condition: weights.type == DataType::kFLOAT || weights.type == DataType::kHALF || weights.type == DataType::kINT32

该错误是在TensorRT-5.1.2.2环境下发生的，错误原因是pytorch模型转onnx模型时某些层的数据类型转为了int64，但是该版本的TensorRT还不支持int64

解决办法升级新版本TensorRT5.1.5.0及其以后版本既可以解决

 

2）、[TensorRT] ERROR: ../rtSafe/cuda/caskConvolutionRunner.cpp (290) - Cask Error in checkCaskExecError<false>: 7 (Cask Convolution execution)
[TensorRT] ERROR: FAILED_EXECUTION: std::exception

这个问题很奇怪，刚开始在google上搜一堆，说是什么构建engine和cuda context不在同一个线程里：下面来自github的答案https://github.com/NVIDIA/TensorRT/issues/301

我不知道我遇到的和他说的本质上是不是一样的，貌似是cuda初始化的问题，我在构建tensorRT engine和context之后使用了img = img.cuda()这段代码就报错，如果在构建tensorRT engine和context之前使用，则是没有问题，下面是代码片段：

engine = onnx_build_engine(onnx_file_path)
inputs, outputs, bindings, stream, context = allocate_buffers(engine)
img = img.cuda() #img是用torch的tensor格式
inputs[0].host = img.cpu().numpy()
OutPuts = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)

上面这样写就报错

img = img.cuda() #img是用torch的tensor格式
engine = onnx_build_engine(onnx_file_path)
inputs, outputs, bindings, stream, context = allocate_buffers(engine)
inputs[0].host = img.cpu().numpy()
OutPuts = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)

这样写就没有问题，似乎要在构建tensorRT engine之前先使cuda初始化一下

3）、在将onnx模型或caffe模型导入RT时，虽然成功了，但是网络的输出结果始终匹配不上，这种情况发生在有某些多分支输出的网络，也不是所有的多分支输出网络会有这样的问题，作者就正好遇到了，也不知道是不是RT版本的bug还是什么，这里7.0.0.11就有这个问题。

当然这样的问题也是有解决办法的，先分析下出现这样问题的可能情况有哪些：

1）加载的onnx模型（或者其他框架的模型）的权重就不对，结果当然不对；

2）权重都保证正确的情况下，某一层的处理方式和其他框架下的处理方式不同，可以设置不同的超参数；

作者通过查看某些中间层的输出来定位问题出现在那一层，发现问题正好在分支融合那里，是个什么情况呢，我大致描述一下：

加载的是一个目标检测网络，有三个分支输出，FPN那样的结构，4个输出；正常情况下是在RT代码里面设置好输入输出层名即可，但是该模型不行，还需要在设置输出层名的队列里面加上三个融合分支的层名，文字可能说得不容易明白，我给个大致的图，如下：

                                                                          

如果514层属于主干的话，517层就属于分支，则需要将517加入到RT的输出层名队列中，同理，其他三个分支也要这样处理。导致这样的问题不知道是不是RT在序列话网络的时候，没把分支给加进去，还是什么，目前还不清楚，但是这样做很简单，也能愉快的解决问题。

       上面那个图是caffe模型的网络可视化图，作者也尝试了，caffe模型添加中间层名到输出层名队列中，是没有毛病的，但是同样的方式用在onnx模型身上就不行了，通过查看onnx模型结构的时候，发现onnx的输入输出在导出onnx模型的时候就固定了，RT加载onnx模型的时候，就不能像加载caffe模型那样输出中间层的结果，那咋办呢？还是有办法的，就给你想要的那个中间层在导出onnx模型时设置一个输出，还是以上图为例子吧，想要输出517层，就给517层再接一个没有权重的层，比如relu，sigmoid，dropout等，再导出onnx模型。反正中间层的输出咋样也不会影响到结果输出，也不会影响到RT的inference推理速度。

    通过使用上面的操作之后，现在的RT输出结果就完全正确了。说了一大堆，其实就是一句话，在设置RT输出的时候，把分支层名也加进去，只是不使用而已。用了这么久的RT，还是第一次遇到这样奇葩的问题，如果你没遇到就万事大吉，遇到了就试试这样的骚操作吧。。。。

4）、在pytorch测试阶段需正确使用with torch.no_grad()和model.eval()，前者可以使测试inference阶段不启用梯度更新，主要是节省显存，不然很容易在批量测试图片时导致显存爆炸；后者是设定网络为测试阶段，并使用网络模型中训练好bn层中的全部权重，不然的话输出结果会有差异。下面是这两句话的官方解释：

Use both. They do different things, and have different scopes.
with torch.no_grad： disables tracking of gradients in autograd.
model.eval()： changes the forward() behaviour of the module it is called upon. eg, it disables dropout and has batch norm use the entire population statistics