TensorRT5.1.5.0 实践 Pytorch2Onnx ，Onnx2TensorRT(python)

TensorRT5.1.5.x的关于onnx的demo只有两个，一个是c++的introductory_parser_samples,一个是python的Object Detection With The ONNX TensorRT Backend In Python。
参考了各方资料，这篇文章是记录+实践
最后成功的环境： Ubuntu 16.06，python 3.5.2，torch 1.0.0，torchvision 0.2.2 onnx 1.4.0 ，TensorRT 5.1.5.0 , OpenCV 2.4.9.1

pytorch转onnx

什么是onnx

Open Neural Network Exchange是开放生态系统的第一笔，它使人工智能开发人员可以在项目的发展过程中选择合适的工具；ONNX为AI models提供了一种开源格式。
它定义了一个可以扩展的计算图模型，同时也定义了内置操作符和标准数据类型。

torch 2 onnx

torch内置方法

torch内置torch.onnx模块，这个模块包含将模型导出到ONNX IR格式的函数。这些模型可以被ONNX库加载，然后将他们转化成可以在其他深度学习框架上运行的模型。
本文就主要记录这个方法，因为官方两个onnx的demo，c++的直接用准备好的onnx文件，python是用.cfg，.weight文件，没有真正的由pytorch代码转。
torch转onnx比较简单，就两句话：

dummy_input=Variable(torch.randn(64,1,28,28,device='cuda'))
output=torch.onnx.export(model,dummy_input,filepath,verbose=True

第一句话设置input格式，第二句话进行export，这个网上资源很多，不详述。

利用torch的cfg和weight，构建onnx的图结构的方法

这个方法就对应着TensorRT官方给的demo，具体请看Object Detection With The ONNX TensorRT Backend In Python.

onnx 2 tensorRT

这里和caffe 2 tensorRT相差不多，遇到了一个问题，转的.onnx有问题，详见“遇到的问题----onnx转换错误”

遇到的问题：

onnx转换错误

按照tensorRT给的API，运行过程中报错

  File "/media/boyun/6a2d0d8c-27e4-4ad2-975c-b24087659438/pycharm/self_pytoch_tensorRT/onnx_tensorRT.py", line 34, in build_engine
    f.write(engine.serialize())
AttributeError: 'NoneType' object has no attribute 'serialize'

提示，文件为none，不能serialize；然后用debug看，用build_cuda_engine创建的engine果然为None：

engine = builder.build_cuda_engine(network)    # engine: None

仔细检查，只有可能是.onnx不太对。
最初，我把.onnx转为可视的模式，从debug中看，发现’‘doc_string’'后面是类似路径的东西，我以为是乱码:

    doc_string: "/usr/local/lib/python3.5/dist-packages/torch/nn/functional.py(396): max_pool2d\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/pooling.py(142): forward\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py(465): _slow_forward\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py(475): __call__\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/container.py(91): forward\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py(465): _slow_forward\n/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py(475): __call__\n/media/boyun/6a2d0d8c-27e4-4ad2-975c-b24087659438/pycharm/self_pytoch_tensorRT/load_model_to_onnx.py(29): forward\n/usr/local/lib/python3

但是，后来用别的.onnx，能够成功生成engine的.onnx，发现其中也有这种“乱码”，所以，其实pytorch转onnx时的.onnx文件都会有这种“乱码”，并不是error。
所以，可以肯定，是我用pytorch转.onnx的时候出了问题。
经过几次排查，排查了.cuda()的问题，排查了是否必须转一步.pt模型格式再转.onnx的问题，最后觉得问题就出在模型结构本身，我跑的是最简单的mnist的三个conv一个dense的卷积神经网络，我把最后的dense（nn.Linear()）删掉，直接让其输出conv3，然后输出.onnx，用这个.onnx再跑转TensorRT，这次engine就没问题了。

engine = builder.build_cuda_engine(network)      #<tensorrt.tensorrt.ICudaEngine object at 0x7fa684965d50>

这就说明TensorRT的onnx解析器，是不支持最后这个层的。于是产生怀疑，难道TensorRT的ONNX解析器连pytorch转onnx的nn.Linear()层都不认识？
后来发现，并不是nn.Linear()的问题，在看了Flatten问题和github上的解释ONNX Export Flatten operator 之后，发现这个view的用法（升维）是不可以的：

res = conv3_out.view(conv3_out.size(0), -1)

要用flatten代替

res=conv3_out.flatten(1)

但是，发现改为flatten之后，pytorch转onnx又报错了

ONNX export failed: Couldn't export operator aten::flatten

然后我就考虑到可能是版本问题，因为我用的是pytorch0.4.0，torch.onnx生成的onnx模块应该是opset 6，但是TensorRT的研发手册上写着需要opset 9，所以更新了torch。
最后，更新完的torch能够识别flatten层，不再报不存在flatten操作符的错，顺利的输出了onnx文件，然后用TensorRT读它，也能顺利读到数据然后顺利创建engine，最后的结果也是正确的。

mnist的结果显示不对

刚开始没有管上面onnx转换错误的问题，直接用TensorRT/data/mnist/mnist.onnx来做的尝试。

do_inference()方法的输出问题

我的代码中，直接用的官方给的TensorRT/sample/python/common.py中定义的do_inference()方法，如下所示

    with get_engine(onnx_path,engine_file_path) as engine,engine.create_execution_context() as context:
        inputs,outputs,bindings,stream=common.allocate_buffers(engine)
        print('Running inference on image ...')
        inputs[0].host=input
        trt_outputs=common.do_inference(context,bindings=bindings,inputs=inputs,outputs=outputs,stream=stream)

通过看debug，我发现推断结果保存在outputs中，下面显示的outputs的debug数据：

outputs = {list} : [Host:\n[  520.9111    -648.3202   -1794.1155      56.337093   133.29428\n  1400.0497    1162.5737   -2256.3643     392.9694    -955.86035 ]\nDevice:\n]
 0 = {HostDeviceMem} Host:\n[  520.9111    -648.3202   -1794.1155      56.337093   133.29428\n  1400.0497    1162.5737   -2256.3643     392.9694    -955.86035 ]\nDevice:\n
 __len__ = {int} 1

即，其实这是个HostDeviceMem类型的数据（其实这里我还没完全清楚），但是肯定不能直接当list处理。（我直接取，确实失败了）
于是参照官方python的demo里的写法，要用trt_outputs取值，而不是outputs，再用debug一看，确实，trt_outputs中保存的是数组结构：

trt_outputs = {list} : [array([  520.9111  ,  -648.3202  , -1794.1155  ,    56.337093,\n         133.29428 ,  1400.0497  ,  1162.5737  , -2256.3643  ,\n         392.9694  ,  -955.86035 ], dtype=float32)]
0 = {ndarray} [  520.9111    -648.3202   -1794.1155      56.337093   133.29428\n  1400.0497    1162.5737   -2256.3643     392.9694    -955.86035 ]
__len__ = {int} 1

所以最后我的后处理代码变成这样：

   prob=[]
   val=0
   sum=0
   for i in range(len(trt_outputs[0])):
       prob.append(np.exp(trt_outputs[0][i]))
       sum+=prob[i]
   for i in range(len(trt_outputs[0])):
       prob[i]/=sum
       val=max(val,prob[i])
               if (val==prob[i]):
           idx=i
   print("The number is %s ,conficence is %s"%(idx,val))

注意trt_outputs相当于（1,10）的array，计数的时候还是走了些坑的。