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} <class '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<pycuda._driver.DeviceAllocation object at 0x7f48259ed620>]
 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<pycuda._driver.DeviceAllocation object at 0x7f48259ed620>
 __len__ = {int} 1

即，其實這是個HostDeviceMem類型的數據（其實這裏我還沒完全清楚），但是肯定不能直接當list處理。（我直接取，確實失敗了）
於是參照官方python的demo裏的寫法，要用trt_outputs取值，而不是outputs，再用debug一看，確實，trt_outputs中保存的是數組結構：

trt_outputs = {list} <class '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，計數的時候還是走了些坑的。