Jetson Orin NX 8GB x JetPack 6.1 x PyTorch

先日は先の手順で、Jetson Orin NX 8GBを搭載したVIA製のAMOS-9100をセットアップしました。今回、新規にPyTorchの実行時間を他のプラットフォームと比較する必要ができたため、JetPack 6.1がセットアップされた環境に 「PyTorch」 をインストールする手順をまとめました。

やや結論

PyTorchでモデルを定義し、ONNXに書き出して、TenosrRTで推論すると爆速。

事前のセットアップ手順

本手順は、以下のセットアップが完了しているJetsonを対象としています。

参考にしたサイト

セットアップ手順

それでは早速、JetPack 6.1のインストールされたJetson Orin NX上にPyTorchをインストールしてみましょう。 ほぼ、参考にしたサイトの手順と同じです。 本記事は、補足としてスクリーンショットとリソースをダウンロードするサイトのURLを埋め込んでいます。

必要なライブラリをインストールし、Python仮想環境をセットアップする

まず、以下の手順でPyTorchの実行に必要な画像処理向けのライブラリ(libjpeg, libpng, libtiff)と、Pythonの仮想環境をセットアップします。

$ sudo apt-get install -y python3-pip libjpeg-dev libpng-dev libtiff-dev
$ sudo apt-get install -y python3.10-venv
$ python3 -m venv ~/env
$ source ~/env/bin/activate

PyTorchのダウンロードしてインストールする

次に、PyTorchをNVIDIAのダウンロードサイトから入手し、インストールします。

$ wget https://developer.download.nvidia.com/compute/redist/jp/v61/pytorch/torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl
$ pip install ./torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl

cuSPARSELtをインストールする

torchvisionをインストールするために、torchvisionが依存するパッケージである 「cuSPARSELt」 をインストールします。

$ cd ~
$ wget https://developer.download.nvidia.com/compute/cusparselt/0.7.1/local_installers/cusparselt-local-tegra-repo-ubuntu2204-0.7.1_1.0-1_arm64.deb
$ sudo dpkg -i cusparselt-local-tegra-repo-ubuntu2204-0.7.1_1.0-1_arm64.deb
$ sudo cp /var/cusparselt-local-tegra-repo-ubuntu2204-0.7.1/cusparselt-*-keyring.gpg /usr/share/keyrings/
$ sudo apt-get update
$ sudo apt-get -y install libcusparselt0 libcusparselt-dev

tochvisionをソースコードからビルドしてインストールする

以下の手順でtorchvisionのソースコードを入手して、Python仮想環境へインストールします。PyTorchの依存関係により、numpy==1.26.4(1.x系)が必要になるため、こちらも併せて更新しましょう。

$ cd ~
$ git clone https://github.com/pytorch/vision.git
$ cd vision
$ git checkout tags/v0.20.0
$ python3 setup.py install
$ pip install -U numpy==1.26.4
$ pip install onnx

PyTorchが実行できることを確認する

以下のPythonスクリプトをJetson上に作成し、実行することでGPU付きのPyTorchが実行できることを確認することができます。

import torch
import torchvision

print("Torch Version:", torch.__version__)
print("TorchVision Version:", torchvision.__version__)
print("CUDA Available:", torch.cuda.is_available())

if torch.cuda.is_available():
print("CUDA Device:", torch.cuda.get_device_name(0))

実行手順は以下の通りです。

$ source ~/env/bin/activate
$ cd ~/sandbox/python/segmentation
$ python check_gpu.py
# Torch Version: 2.5.0a0+872d972e41.nv24.08
# TorchVision Version: 0.20.0a0+afc54f7
# CUDA Available: True
# CUDA Device: Orin

記事は以上となります。手順を忘れそうだったため、メモとして本記事を投稿させていただきました。最近PyTorch向けのソースコードを書いていないので、そろそろPyTorch周りでトライしたこともまとめていきたいと思います。

次回の記事もご期待ください！