Zenn
👋

TensorRTを試してみる - Python環境構築 -

2020/10/23に公開
ディープラーニング
IoT
tech

はじめに

過去2つの投稿を基にJetsonNano上でのYOLOv4-tinyをTensorRTで動かす実践をしてみます。
Jetson NanoでDarknetのYolov4を試してみる
TensorRTを試してみる - TensorRTとは -

個人的に様々な環境で取り扱いが容易なYOLOv4-tinyのONNXを作り上げ、
そこからTensorRTへのImportを実行して動かしてみることを狙います。

今回やりたいことを項目で上げていくと、

  • Yolov4-tiny ONNX Modelを用意
  • TensorRTでONNX Modelを取り込み
  • 下記のそれぞれについて単一画像でのmAPを確認
  • YOLOv4-tiny by darknet
  • YOLOv4-tiny by TensorRT
  • YOLOv4-tiny by TensorRT(FP16)
  • 下記のそれぞれについてカメラ入力でのFPSを確認
  • YOLOv4-tiny by darknet
  • YOLOv4-tiny by TensorRT
  • YOLOv4-tiny by TensorRT(FP16)

一応公式実装もあるのですが、自前で実装を試みてみます。

なお、JetsonNano内にPythonでの環境を整えること自体に手こずったため、
本記事ではPythonでの環境構築に関してまとめます。

ONNX & Pytorch Setup

まずはImportするためのONNX Modelを作りますが、darknetから直接ONNXにする手段はないので、
この部分だけは他者の力をお借りします。pytorch-YOLOv4を利用します。

python3環境が整ってなかったのでまずはpipの取り込み&upgrade実行

$ sudo apt-get --reinstall install python3-setuptools python3-wheel python3-pip
$ sudo -H pip3 install --upgrade pip

python3でonnxを使う必要があるようなのでまずはonnxのセットアップ。
pip3実行時にWARNINGが出るようなので、python3経由でpip実行。

$ python3 -m pip install --user cython protobuf numpy
$ sudo apt-get install libprotobuf-dev protobuf-compiler
$ python3 -m pip install --user onnx

まずは必要なPython Packageを取り込み...

$ python3 -m pip install --user -r requirement.txt
    ...
ERROR: Could not find a version that satisfies the requirement torch==1.4.0 (from -r ./requirements.txt (line 2)) (from versions: 0.1.2, 0.1.2.post1, 0.1.2.post2)
ERROR: No matching distribution found for torch==1.4.0 (from -r ./requirements.txt (line 2))

ダメみたいですね...おそらくですがARM64 Linuxな影響でpython系のpackageが全体的にnot supportになっていると考えました。
いくつか情報を漁ってみて試してみました。Dockerfileがあったのでこちらを参考にインストールを進めて動かしてみればいけるんじゃないかと予想。
なお、CUDA 10.2がtorch 1.5.0からしか対応してない上、Pytorch1.5.0はJetPack4.4 Production Releaseには非対応なので、1.6.0を導入。
おそらくgitリポジトリそのままでは動かせなさそうなので、元リポジトリをforkして作業を進める。

$ sudo apt install -y --no-install-recommends python3-pip python3-dev libopenblas-dev libopenmpi2 openmpi-bin openmpi-common gfortran
$ wget --quiet --show-progress --progress=bar:force:noscroll --no-check-certificate https://nvidia.box.com/shared/static/yr6sjswn25z7oankw8zy1roow9cy5ur1.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
$ python3 -m pip install --user torch-1.6.0-cp36-cp36m-linux_aarch64.whl --verbose
$ rm torch-1.6.0-cp36-cp36m-linux_aarch64.whl

$ sudo apt update && sudo apt install libjpeg-dev
$ git clone -b v0.7.0 https://github.com/pytorch/vision torchvision
$ cd torchvision
$ sudo python3 setup.py install
$ sudo apt update && sudo apt install -y --no-install-recommends sox libsox-dev libsox-fmt-all
$ cd ..

# 多分要らない
$ git clone -b v0.6.0 https://github.com/pytorch/audio torchaudio
$ cd torchaudio
$ sudo python3 setup.py install
$ cd ..

$ export PATH=/usr/local/cuda/bin:${PATH}
$ export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}
$ python3 -m pip install --user pycuda --verbose

上記まで対応させ、一通りimportが通ることを確認完了!

Pytorch YOLOv4の動作確認

まずはforkしたリポジトリをcloneし、必要なパッケージを取得。(先ほど取得したパッケージは取得しないようにrequirement.txtに変更)

$ git clone https://github.com/sa-kei728/pytorch-YOLOv4 && cd pytorch-YOLOv4
$ sudo apt update
$ python3 -m pip install --user -r ./requirements.txt

また、OpenCVを導入する必要があるので、このタイミングで導入します。
当初はpipで取得を試みたのですが、aarch64向けのものはpipで取得できないようだったので、
先駆者のスクリプトを用いてOpenCV4.3.0を導入。
1時間半ほどかかります。

$ cd install_opencv
$ ./run_install_opencv.sh

その後、pytorch-YOLOv4が動作するかどうかを確かめる。
Userモードでカメラを動かしたい都合上、下記でUSERをvideo groupに登録しておく。
※USERは各々の設定に置き換えてください。

$ usermod -aG video USER

yolov4-tinyのテストを行う。
※予めyolov4-tiny.weightを移しておく。

$ python3 demo.py -cfgfile cfg/yolov4-tiny.cfg -weightfile yolov4-tiny.weights

下記のように実行結果が出ており問題がなさそうなことを確認できた。
torch_yolov4_tiny

※なお、imgfileを指定して推論を動かすとSegmentation faultが発生してしまうが、
gdbを見たところ、cudnn絡みの問題なのとプロセス終了時に発生するので問題解決は諦めました。

ONNX動作確認

ONNX部分も動作するかどうかを確認してみました。

$ python3 demo_darknet2onnx.py cfg/yolov4-tiny.cfg yolov4-tiny.weights data/dog.jpg 1

predictions_onnx.jpg

おわりに

Pythonでの環境構築の沼にハマってしまったので今回は本質部分まで辿れませんでしたが、
ある程度まともな環境が出来たため、次回は調査を進めていきたいと思います。

Discussion