Jetson Orin NX 8GB [JetPack 6.1] にPyTorch(torch, torchvision) をインストールする

# Jetson Orin NX 8GB x JetPack 6.1 x PyTorch

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

#### やや結論

<font color="blue">**PyTorchでモデルを定義し、ONNXに書き出して、TenosrRTで推論すると爆速。**</font>

#### 事前のセットアップ手順

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

https://qiita.com/xtrizeShino/items/4227a60108caee350042

#### 参考にしたサイト

https://github.com/azimjaan21/jetpack-6.1-pytorch-torchvision-

## セットアップ手順

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

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

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

```bash
$ 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のダウンロードサイトから入手し、インストールします。

https://developer.nvidia.com/embedded/downloads#?search=PyTorch

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3919880/894719c1-8084-4b48-9269-7c8ea8cfca5e.png)

```bash
$ 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」** をインストールします。

https://developer.nvidia.com/cusparselt-downloads?target_os=Linux&target_arch=aarch64-jetson&Compilation=Native&Distribution=Ubuntu&target_version=22.04&target_type=deb_local

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3919880/123142d9-5c59-4f86-93ad-6b8bf343bbbb.png)

```bash
$ 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系)が必要になるため、こちらも併せて更新しましょう。

```bash
$ 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が実行できることを確認することができます。

```python
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))
```

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

```bash
$ 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周りでトライしたこともまとめていきたいと思います。

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