🆕

Jetson Orin nanoにpytorch+torchvisionをインストールする

2023/11/27に公開

PyTorch

jetson orin nanoで画像処理したいけど、yoloとか使いたいので環境構築した。
pytorchとtorchvisionをバージョンに気をつけてインストールしないといけない。
その備忘録。

使用環境

Jetson Orin Nano Dev Kit

JetPack SDK v5.1.2

忙しい人のために

# いつものやつ
sudo apt update
sudo apt upgrade
# jetpackをインストール
sudo apt install nvidia-jetpack
# pytorchをインストール
sudo apt install autoconf bc build-essential g++-8 gcc-8 clang-8 lld-8 gettext-base gfortran-8 iputils-ping libbz2-dev libc++-dev libcgal-dev libffi-dev libfreetype6-dev libhdf5-dev libjpeg-dev liblzma-dev libncurses5-dev libncursesw5-dev libpng-dev libreadline-dev libssl-dev libsqlite3-dev libxml2-dev libxslt-dev locales moreutils openssl python-openssl rsync scons python3-pip libopenblas-dev
## 下記は自分のインストールしたいバージョンを指定、jetpackのバージョンに依存
export TORCH_INSTALL=https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
python3 -m pip install --upgrade pip
python3 -m pip install aiohttp numpy=='1.19.4' scipy=='1.5.3'
export "LD_LIBRARY_PATH=/usr/lib/llvm-8/lib:$LD_LIBRARY_PATH"
python3 -m pip install --upgrade protobuf
python3 -m pip install --no-cache $TORCH_INSTALL
# torchvisionをビルド＋インストール
sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
インストールしたい任意の場所にcd
git clone --branch release/0.16 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.16.1  # where 0.x.0 is the torchvision version  
python3 setup.py install --user

セットアップ

下記URLに従ってjetson orin nanoをセットアップします。

Jetpack SDKをインストールする

前項の初期設定を実行するだけでJetpackはインストールされていると思います。なので多くの人はこの項を読み飛ばしてもらって大丈夫です。

JetPackとは、NVIDIAが提供するソフトウェア開発キット（SDK）です。
Jetsonプラットフォーム向けに、主にエッジAIやロボティクスで使用するツール・API・ライブラリを詰め込んでくれてあります。
これをインストールするだけで、CUDAやcuDNNやTensorRTなどを手動でインストールせず、迅速に開発を始めることができます。

弊環境ではNVMeスロットに挿したSSDにJetsonイメージをフラッシュしており、^[1] ^[2]
参照した作業ではJetpackがインストールされていないようだったので手動でインストールしました。

Jetpackをインストール

sudo apt update
sudo apt install nvidia-jetpack

手動でインストールしたためか、cudaやnvccがうまく読み込めていませんでした。

.bashrcに下記を追記してPATH環境変数に追加

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

nvccが読み込めているか確認

nvcc -V

なお、SDK ManagerというNVIDIAのツールを使って外部のUbuntu PCからjetsonイメージをjetson boardにフラッシュした場合も、jetpack相当（CUDAやcuDNNなど）がインストールされるのでこの項の実行は不要です。

Pytorchをインストールする

JetsonはARMアーキテクチャのチップを搭載しており、intel core-iシリーズなどのx86_64アーキテクチャ用にビルドされた通常のpytorchが動作しません。

多くのpipでインストールするパッケージはpythonで実装されているのでCPUアーキテクチャにほとんど依存しませんが、一部の複雑なパッケージはCPUアーキテクチャに対応したパッケージを使う必要があります。

これを解決するため、NVIDIA公式から事前にビルドされたpip wheelファイルが配布されています。Jetsonのユーザーはこれでpytorchをインストールするのが良いです。

インストールするべきバージョンはJetpack SDKのバージョンに依存します。
ここ^[3] にアクセスすると任意のjetpack（遷移先URLのv512はjetpack v5.1.2を意味する。他のバージョンも参照できる。）に対応したバージョンのpytorchのpip wheelファイルを参照できます。

依存関係の解決

sudo apt update
sudo apt install autoconf bc build-essential g++-8 gcc-8 clang-8 lld-8 gettext-base gfortran-8 iputils-ping libbz2-dev libc++-dev libcgal-dev libffi-dev libfreetype6-dev libhdf5-dev libjpeg-dev liblzma-dev libncurses5-dev libncursesw5-dev libpng-dev libreadline-dev libssl-dev libsqlite3-dev libxml2-dev libxslt-dev locales moreutils openssl python-openssl rsync scons python3-pip libopenblas-dev

インストールするpytorchを環境変数で指定
環境変数に自身のjetpackのバージョンに対応するpip wheelファイルのURLを設定する。
(先ほどのアクセス先からURLをコピペ)

export TORCH_INSTALL=https://developer.download.nvidia.com/compute/redist/jp/v512/pytorch/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

pytorchをインストール。

python3 -m pip install --upgrade pip
python3 -m pip install aiohttp numpy=='1.19.4' scipy=='1.5.3'
export "LD_LIBRARY_PATH=/usr/lib/llvm-8/lib:$LD_LIBRARY_PATH"
python3 -m pip install --upgrade protobuf
python3 -m pip install --no-cache $TORCH_INSTALL

torchvisionをビルド＋インストールする

pytorchのバージョンと対応するバージョンのtorchvisionをインストールする必要がある。
https://github.com/pytorch/vision
この記事の冒頭のリンク先にその依存関係の対応表が載っている。

先ほどインストールしたpytorch v2.1.0に対応するtorchvisionのバージョンは記載がないが、github上のtorchvisionのREADME^[4] を参照すると、v0.16系が対応している模様。リリースノートを参照するとバグが修正されたv0.16.1が最新（2023/11/27時点）なのでこれをビルド+インストールする。

依存関係の解決

torchvisionをビルドするために必要な画像処理やビデオ処理に関する依存関係のライブラリをインストールします。

sudo apt install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
sudo apt install ffmpeg libavutil-dev libavcodec-dev libavformat-dev libavdevice-dev libavfilter-dev libswscale-dev libswresample-dev libswresample-dev libpostproc-dev libjpeg-dev libpng-dev libopenblas-base libopenmpi-dev
python3 -m pip install ninja ffmpeg pynvjpeg

ビルドのための環境変数設定

torchvisionのビルドには、CUDAやcuDNNなどのNVIDIAのライブラリを利用します。
これを有効にするために、環境変数を設定します。

export USE_CUDA=1 USE_CUDNN=1 USE_MKLDNN=1 TORCH_CUDA_ARCH_LIST="8.7" FORCE_CUDA=1 FORCE_MPS=1

torchvisionをビルド＋インストール

任意のディレクトリに移動し、GitHubからtorchvisionのリポジトリをクローンします。

cd 任意のディレクトリ
git clone --branch release/0.16 https://github.com/pytorch/vision torchvision
cd torchvision

Pythonのセットアップツールを使ってtorchvisionをインストールします

python3 setup.py install --user

以上です！

なお基本的に以下を参照しています。誰かのお役に立てば幸いです。

脚注