我的 Jetson Nano 是安裝 Jetpack 4.4.1 所以 Cuda 版本是 10.2，我已經安裝好 torch 1.6.0 和 torchvision 0.7.0，若還沒安裝的可以回去參考使用 PyTorch Installer 安裝。
先安裝一些相依套件，並下載 nanodet 的 repo
$ pip3 install Cython termcolor numpy tensorboard pycocotools matplotlib pyaml opencv-python tqdm
$ git clone https://github.com/RangiLyu/nanodet.git
$ cd nanodet
$ python setup.py develop下載 nanodet 的 coco pretrained weight
$ pip3 install gdown
$ gdown --id '1EhMqGozKfqEfw8y9ftbi1jhYu86XoW62'將 pth 轉成 onnx，運行 export.py 可以得到 output.onnx
!python3 tools/export.py原本的 pth 網路架構輸出層是輸出 3 個 scale：10x10, 20x20, 40x40
- class: [1, 80, 10, 10], [1, 80, 20, 20], [1, 80, 40, 40]
- bbox: [1, 32, 10, 10], [1, 32, 20, 20], [1, 32, 40, 40]
轉完的 onnx 會經過 sigmoid 和 reshape 變成 (在 pytorch 中在後處理時會經過這些步驟)
- class: [1, 100, 80], [1, 400, 80], [1, 1600, 80]
- bbox: [1, 100, 32], [1, 400, 32], [1, 1600, 32]
我們可以將轉好的 output.onnx 放到 netron 查看網路架構
為了方便起見我們直接更改 outputs 的位置，我們使用 onnx-utils 來更改，--outpus 第一個位置放 output name，[]裡放 output shape，例如 831[1,80,10,10]，點選 node 可以查看資訊。
$ wget https://raw.githubusercontent.com/saurabh-shandilya/onnx-utils/master/onnx_edit.py
$ python3 onnx_edit.py output.onnx nanodet.onnx\
--outputs '787[1,80,40,40], 788[1,32,40,40], 809[1,80,20,20], 810[1,32,20,20], 831[1,80,10,10], 832[1,32,10,10]'更改後的網路架構如下
將更改完的 nanodet.onnx 轉成 tensorrt 的 engine，如何將 onnx 轉成 trt 可以參考以下文章
$ wget https://raw.githubusercontent.com/d246810g2000/tensorrt/main/build_engine.py
$ python3 build_engine.py -o nanodet.onnx -t nanodet.trt接著我們就可以來 inference 了
先來試試作者提供的 demo.py
$ python3 demo/demo.py csicam --config config/nanodet-m.yml --model nanodet_m.pth可以看到 FPS 大概在 7 左右，模型推理時間為 80 ms，解碼時間為 40 ms
再來試試用 tensorrt 優化後的 trt_demo.py，為了方便起見我直接將 nanodet.trt 的輸出轉換成 torch.tensor 再進行後處理，結果如下
# trt_demo.py 需要 import common
$ wget https://raw.githubusercontent.com/d246810g2000/tensorrt/main/nanodet/common.py
$ wget https://raw.githubusercontent.com/d246810g2000/tensorrt/main/nanodet/trt_demo.py
$ python3 demo/trt_demo.py csicam --config config/nanodet-m.yml --model nanodet.trt可以看到 FPS 大概在 9 左右，模型推理時間從 78ms 提升至 25ms 左右，可以看到 TensorRT 加速了大概 3 倍左右，可以看到是 decode 的部分大大的影響了 FPS 的數字
