Jetson Orin Nano開発者Kitにて、TensorRT使ってカスタムモデルを動かすためにしたセットアップ

こんにちは、ピリカ開発チームの九鬼です。

ごみ検出モデルをエッジAIで動かすにあたり、Jetson Orin Nano開発者キットでいくつかセットアップしました。その中でトラブルのあったポイントがいくつかあったので、その内容とセットアップ方法を共有いたします。

前提

SDK Managerは使わず、Imageを直接書き込む方法でJetpackをセットアップしています。

1. OSセットアップ: UEFI Firmwareを最新にする

工場出荷時、機器によってはUEFI Firmwareが36.3未満で古い場合があります。この環境では最新のJetpack 6.xを起動できません1。この場合Firmwareのアップデートが必要です。

1.1. UEFI Firmwareのバージョンを確認する

SDカードなしで起動し、ブート画面でESCを押します。するとBIOSの設定画面が表示され、左上にてバージョンを確認できます。下図の場合、3.0-32616947です。32.6よりも古いので、アップデートが必要です。

Jetson Orin Nano 開発者キット BIOSでのUEFI Firmwareバージョン表示
Jetson Orin Nano 開発者キット BIOSでのUEFI Firmwareバージョン表示

1.2 UEFI Firmwareをアップデートする

以下の3ステップでアップデートします。詳しい手順は公式マニュアルのUpgrade the Jetson UEFI firmware to 36.xに従ってください。

  • Jetpack 5.1.3をSDカードに用意し、起動する: バージョンが5.0-35550185に更新されます。
  • Jetpack 5.1.3上でQSPI Updater packageをインストールし、再起動する: バージョンが36.3.0に更新されます。
    • この時点でJetpack 5.1.3での起動ができなくなります。

あとはJetpack 6.1をSDカードに用意すると、Jetpack 6.1で起動できます。

2. Wi-Fiを設定する

始めからWi-Fiモジュールが入っているため、ホットスポットを指定すればそのままWi-Fiに接続できます。

2.1. GUIから設定する

ubuntuのネットワークマネージャが使用できます。設定 > Wi-FiからSSID・パスワードを設定すればOKです。

2.2. CUIから設定する

nmcliを使って設定します。

例えば、設定名が"wifi-con", SSIDが"SSID-1234"とすると以下のようになります。

  • ネットワーク設定追加2
    • sudo nmcli connection add type wifi ifname wlan0 con-name "wifi-con" ssid "SSID-1234" wifi-sec.key-mgmt wpa-psk wifi-sec.psk "つなぎたいホットスポットのパスワード"
  • ネットワークに接続
    • sudo nmcli connection up "接続名"

3. デスクトップGUI無効化/有効化

機器を自律稼働させるとき、GUIが不要になります。以下のコマンドで無効にできます。

sudo systemctl set-default multi-user.target

反対に、GUIが必要になった場合は以下のコマンドで有効にできます。

sudo systemctl set-default graphical.target

4. TensorRTにて、GPUを使って推論する

4.1. ONNX形式のモデルを、TensorRTで動作できるように最適化・変換する

Jetpack 6では、デフォルトでTensorRTが導入されています。この中にtrtexecバイナリがあり、本バイナリで変換できます。ただし、デフォルトではパスが通っていないので、パスを通します。

$ export PATH=/usr/src/tensorrt/bin/:$PATH

パスを通したあと、以下コマンドで変換します3

$ trtexec --onnx=model.onnx --saveEngine=model.trt

変換が成功すると、saveEngineで指定した場所にTRT形式のモデルが出力されています。

なお、この変換で使ったTensorRTのバージョンは、推論するときに使うTensorRTのバージョンと一致している必要があります。もしバージョンが合っていない場合、推論できません。

4.2. Python環境の用意

Jetpack 6.1では、Python 3.10.12が入っています。このPython 3.10.12を使って推論するよう、環境を整えます。 CUDAを動かすためのWrapperとしてPyCUDAが提供されているので、こちらをインストールします。

PyCUDAのインストールではビルドが走るので、インクルードパスやライブラリパス一式を入れておきます。

$ export PATH=/usr/local/cuda/bin:$PATH
$ export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
$ export CPATH=/usr/local/cuda/include:$CPATH

その後、PyCUDAをインストールします。正常にビルドできていれば、数分でインストールが完了します。

$ pip install pycuda 

もしcannot find -lcurandなどの表示がされてビルドできない場合、前述のパスが通っていない可能性があります。パスを確認の上、もし通っていなければ上記パス追加を再実行します。

4.3. PyCUDA周りを使っての推論処理の実装

公式サンプル4にしたがって実装します。推論時はGPU上のメモリに入力・出力データの一時保存場所を用意する必要があります。したがって、allocate_memoryでGPU上のメモリを確保しています。推論時には、入力データをGPU上のメモリに書き込み、推論後に出力データをGPU上のメモリから読み出しています。

注意点として、推論をmainスレッド以外で実行しているときは別途GPUバイスのcontextを取得しておく必要があります。もし取得していない場合、LogicError: ... invalid device contextにてコードがクラッシュします。

モデル読み出しと合わせて

import pycuda.driver as cuda

# 中略
cfx = cuda.Device(0).make_context()

等でcontextを得た上で、推論時に

cfx.push()
# 推論処理
cfx.pop()

とデバイスのcontextを制御する必要があります。また、実行を終了するときは当該contextを開放しておく必要があります(開放しなければ当該コードがクラッシュします)。

cfx.pop()
cfx = None

4.4. 推論コードの実行

当該コードでのみ使うライブラリが存在する場合、仮想環境上を作っておきます。PyCUDAおよびTensorRTはグローバルなpipにインストールされているので、system-site-packagesへのアクセスを付与します。

$ python -m venv .venv --system-site-packages

requirements.txtからコード実行に必要なライブラリをインストールします。

$ source .venv/bin/activate
(.venv) $ pip install -r requirements.txt

ライブラリインストール後、コードを実行します。

4.5. TensorRTの推論でのトラブルシューティング

4.5.1. deserializeCudaEngine failed. Serialization assertion magicTagRead == kMAGIC_TAG failed.Magic tag does not match でクラッシュした

モデルの変換に使ったTensorRTと、デシリアライズするTensorRTでバージョンが合っていない場合に発生します。同じJetson上でモデルを変換・推論している限りは問題ありません。

また、単にTRT形式以外のモデルを指定している場合も本エラーになります。そのため、ロードするモデルが合っているか確認すると良いでしょう。

ValueError: ndarray is not contiguous でクラッシュした

入力データをCUDA側にcopyせず、直接参照渡した場合などで発生します。この場合、cuda.pagelocked_emptyで確保した入力データ領域に対し、np.copytoで入力データをコピーして渡すと解消されます。

推論した結果がすべて0出力になる

入力データの渡し忘れ、もしくは出力データの取り出し忘れのいずれかで発生し得ます。この場合、以下を確認します。

  • cuda.memcpy_htod_async で書き込んだ入力において、ホスト側の入力データ領域に入力データをコピーできているか確認する。
  • cuda.memcpy_dtoh_async で書き込んだ出力において、ホスト側の出力データ領域に出力結果をコピーできているか確認する。

a. その他

a.1. exFAT対応

Jetpack 6では、デフォルトでexFATが認識されません。幸いFUSEがモジュールとしてビルド済みのため、以下の手順で有効化できます。

$ sudo apt-get install exfat-fuse
$ sudo ln -s /usr/sbin/mount.exfat-fuse /sbin/mount.exfat

mountは/usr/sbin以下を見に行かないので、シンボリックリンク/sbin下におく必要があります。

a.2. 外部ストレージを自動マウントする

/etc/udev/rules.d/99-nv-ufs-mount.rulesを使うと自動マウントできます。特にUSBメモリなどを使いたいときに有効です。

以下、設定例です。各外部ストレージに対し、/media下に個別のUUIDで自動マウントします。

# Mount UFS card when detected.
ACTION=="add", KERNEL=="sda[1-9]", ENV{ID_FS_UUID}!="", SUBSYSTEM=="block", RUN{program}+="/usr/bin/systemd-mount --no-block --automount=yes --collect $devnode /media/%E{ID_FS_UUID}"

# Unmount UFS card when removed.
ACTION=="remove", KERNEL=="sda[1-9]", ENV{ID_FS_UUID}!="", SUBSYSTEM=="block", RUN{program}+="/usr/bin/systemd-umount /media/%E{ID_FS_UUID}"

fstabを使っても自動マウントできますが、外部SSDなど起動に時間がかかるストレージの場合マウントできない場合があります。自動マウント実行時にストレージが認識されていない場合によく発生します。

a.3. IPアドレスを固定する

閉域網でP2Pにて通信したい場合があります。その場合IPアドレスを固定することで対処できます。

Jetpack 6.1では、netplanを使って設定します。/etc/network/interfacesに記載する方法は効かないのでご留意ください。

最初に、netplanをインストールします。

$ sudo apt-get install netplan.io

その後、設定ファイル(/etc/netplan/00-installer-config.yaml)にIPアドレスを使う旨を記述します。

以下は、IPアドレスを192.168.1.10に固定し、なおかつサブネットを192.168.1.*に限定する場合の設定です。

network:
  version: 2
  renderer: NetworkManager
  ethernets:
    enP8p1s0:  # dmesg | grep eth0 したときに表示される、rename後のイーサネットインタフェース名
      addresses:
      - 192.168.1.10/24
      routes:
      - to: default
        via: 192.168.1.255
      nameservers:
        addresses:
        - 8.8.8.8

  1. Firmwareが対応していないJetpackバージョンで起動しようとすると、"Test Key Detected"と表示された上でUEFIのシェルが表示されて起動が止まります
  2. 機器によってはwlan0でない場合があります。この場合、dmesg上でwlan0をgrepするとリネーム先を確認できます。
  3. 最適化のレベルは引数で調整できますが、ここでは割愛しています。詳しくは公式sampleを参照ください。
  4. onnx_helperという名前ですが、ONNXが関係するのはONNX→TRTへの変換処理のみです。