Skip to content

JetsonNano trt_pose

Eiichiro Momma edited this page Apr 17, 2021 · 2 revisions

(JetsonNano) trt_pose

NVIDIA謹製のボーン抽出。

  • ほぼ配布元に書いてある通りで動くが幾つか注意が必要。

JetPackのインストール

バージョン合わせが命なので注意

Jetpackのサイトのアーカイブから4.5(4.5.1ではない)のイメージを拾ってきてbalenaEtcherとかでSD microカードに焼いてから,差し込んで起動。暫くすると利用規約画面になるので普通に勧める。言語についてはトラブルの原因なので基本Englihで進めていく。

ユーザーは普通に作る。パスワードについては,キーを英語で設定した場合,日本語キーボードを使っていると記号を含めたときにハマるので注意。 領域拡張等の質問にデフォルトで答えて行くと,最後に再起動がかかって作成したユーザーでのログイン画面になる。

ログイン後,時計がおかしい場合はまず合わせる(ntpが塞がれていると起きる)。やらないとWebページでエラーが出る。歯車マークの設定アプリでBrightness&Lockを開き,最低限LockをOFFにしておく。5分でディスプレイの電源を切りにくるのでNeverにしておく。

PyTorchとtorchvisionのインストール

画面左上のUbuntuマークをクリックしてTerminalを起動する。

まず,JetPackのswapがデフォルトで2GBで足りないので変更する。ここのスクリプトを使う。 ターミナル上で下記を実行する。

cd
git clone https://github.com/JetsonHacksNano/resizeSwapMemory.git
cd resizeSwapMemory
./setSwapMemorySize -g 4

で再起動する。

次にJetCamを入れる。入れ方は書いてある通り

git clone https://github.com/NVIDIA-AI-IOT/jetcam
cd jetcam
sudo python3 setup.py install

続けてNVIDIAの用意したpytorchを入れる。 NVIDIAのフォーラムののInstallationの矢印を開き,Python3.6のコマンドを一行ずつTerminalにコピペして実行する。JetPack4.5の場合は1.8なので書かれてる通りで良い。 (sudoで始まるコマンドは管理者権限が必要なので自分のパスワード入力を求められる) 途中死んだかのように思えるが,ビルドしながら入れてるので結構時間がかかる。

wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
pip3 install Cython
pip3 install numpy torch-1.8.0-cp36-cp36m-linux_aarch64.wh

続けてtorchvisionを入れる。これも貼り付けて行くのだが,2行目はバージョンを自分で入れる必要がある。PyTorch1.8に対応したv0.9.0をのところに書いて実行。 4行目は

export BUILD_VERSION=0.9.0

になる。 で,5行目でまたビルドするので待ち。7行目のpipはpip3の間違い。ということで下記のようになる。

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch v0.9.0 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.9.0
python3 setup.py install --user
cd ../  # attempting to load torchvision from build dir will result in import error
pip3 install 'pillow<7'

再びtrt_poseのサイトに戻って"2. install"の通りtorch2trtを入れる。下の3行を順次Terminalにコピペ。

cd
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
sudo python3 setup.py install --plugins

"3. Install other..."も2行をコピペ

sudo pip3 install tqdm cython pycocotools
sudo apt-get install python3-matplotlib

Setp 2もホームディレクトリに移動してから,続けて実行。3行目がビルドなので時間がかかる。

cd ~/
git clone https://github.com/NVIDIA-AI-IOT/trt_pose
cd trt_pose
sudo python3 setup.py install

Step 3で表のWeightsファイルをresnet18とdensenet121の両方をダウンロードし,ファイルマネージャでtrt_poseのtasksのhuman_poseディレクトリに移動する。

次にこのファイルをダウンロードする。(Downloadsフォルダに保存する)

この作業は本来要らない筈だが,Notebook上でダウンロードしようとするとコケるため。

ブラウザを開いていたら一旦閉じる。(chromeがsandboxを使っていてGPUにアクセスできないため) terminalで

cp ~/Downloads/resnet18-5c106cde.pth ~/.cache/torch/hub/checkpoints/
sudo apt install python3-notebook jupyter-notebook
chromium-browser --no-sandbox &
jupyter notebook

で,trt_pose/tasks/human_poseのlive_demo.ipynbを開き,あとはShift+Enterで順次実行していく。↓のセルとその下のtorch.saveのセルは1度実行すれば不要。(時間がかかる)

import torch2trt

model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)

OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'

torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)

はコメントアウトしてしまい,

# import torch2trt

# model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)
# OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'

# torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)

としておけば良い。

次回以降使うときのために

OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'

の行だけ次のセルにコピーしておく

from torch2trt import TRTModule
OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'
model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))

のようになる。

カメラについてはUSBカメラを使うときは

from jetcam.usb_camera import USBCamera
# from jetcam.csi_camera import CSICamera
from jetcam.utils import bgr8_to_jpeg

camera = USBCamera(width=WIDTH, height=HEIGHT, capture_fps=30)
# camera = CSICamera(width=WIDTH, height=HEIGHT, capture_fps=30)

camera.running = True

のまま。ただし,USBカメラしか繋がってなく,video0の状態が前提で,ZIFケーブルのRaspberry piカメラ(CSIカメラ)が既に繋っているときはUSBカメラがvideo1になるので,

from jetcam.usb_camera import USBCamera
# from jetcam.csi_camera import CSICamera
from jetcam.utils import bgr8_to_jpeg

camera = USBCamera(width=WIDTH, height=HEIGHT, capture_fps=30, capture_device=1)
# camera = CSICamera(width=WIDTH, height=HEIGHT, capture_fps=30)

camera.running = True

としてcapture_deviceで番号を指定する必要がある。 CSIカメラを使うときはコメントを外してcameraの差し替え

# from jetcam.usb_camera import USBCamera
from jetcam.csi_camera import CSICamera
from jetcam.utils import bgr8_to_jpeg

# camera = USBCamera(width=WIDTH, height=HEIGHT, capture_fps=30)
camera = CSICamera(width=WIDTH, height=HEIGHT, capture_fps=30)

camera.running = True

で,暫く待っているとgstreamerに画像が流れているようなメッセージがターミナルに出てくるので,エラーになってないか確認。そこのOutput Stream W= ???? H=????がカメラからの画像サイズを意味する。

※CSIカメラのとき,ここで慌てて進もうとするとエラーになるケースがあるのでGstreamerのログをちゃんと見る。エラーがドバーっと出たら取り敢えずJetsonを再起動して再開方法に従って実行すれば良い。

ここから先はipywdgetsのImageでNotebook内に画像を常時表示し,その画像を別プロセスで更新していく感じになる。

import ipywidgets
from IPython.display import display

image_w = ipywidgets.Image(format='jpeg')

display(image_w)

は,実行すると壊れた画像アイコンが表示されるが,次以降のプロセスで更新されるのでそのままにする。

なお,最初の画像サイズが不定のため横長になってしまうので,先程のW=???? H=????の値のアスペクト比を守りつつ程良い大きさを指定してやると良い。例えばW=1280, H=720のストリームだった場合

import ipywidgets
from IPython.display import display

image_w = ipywidgets.Image(format='jpeg', width=640, height=360)

display(image_w)

のように書いておくと,良い感じに表示される。

camera.observe(execute, names='value')

のところまで順次実行するとちゃんと表示されているのがわかる。

止めるときは最後の

camera.unobserve_all()

を実行する。

そうすると先程作ったwidgetにカメラ画像からボーン抽出した結果が表示される。

終了する前にノートブックは保存しておく。

再開方法

Jetson nanoの電源を入れ直した場合

  1. sandboxなしでchromiumを起動
  2. jupyter notebookを起動
  3. live_demo.ipynbを開いて順次実行 とすれば再度利用できる。
Clone this wiki locally