Jetson Orin NX에 torch 2.1.0와 torchvision 0.16.1 설치하기

March 22, 2024, 11:54 a.m.

우연한 기회에 Jetson Orin NX를 만져볼 기회가 생겨 간단하게 yolov5를 돌려보고 싶었는데... 그냥 평상시처럼 torch와 torchvision을 설치하고 실행해보니 CUDA가 잡히지 않는 경우, 혹은 아래와 같은 에러메시지를 볼 수 있었다.

UserWarning: Failed to load image Python extension: '/home/thmos/anaconda3/envs/env1/lib/python3.8/site-
packages/torchvision/image.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKSs’If you don’t plan 
on using image functionality from torchvision.io 2, you can ignore this warning. Otherwise, there might be 
something wrong with your environment. Did you have libjpeg or libpng installed before building torchvision 
from source?

혹은 이와 비슷한 다양한 에러 메시지를 볼 수 있었는데, 구글을 뒤져가며 거의 10시간 만에 해결했다.

시작에 앞서, Jetson 플랫폼에 대해서 알아야 할 것이 몇가지 있다.

1. 기본적으로 Jetson에서는 Python 3.8을 사용하는것을 추천한다.
2. Jetson은 tegra라는 gpu 아키텍처를 사용하기 때문에 기존에 Nvidia GPU를 사용하듯이 torch를 설치하면 작동하지 않는다. -> 전용으로 빌드된 torch를 사용하자.
3. Jetson은 Jetpack이라는 개발 키트가 설치되어있는데, Jetpack의 버전마다 Cuda와 빌드 툴의 버전들이 다르다. -> 이에 맞는 버전으로 빌드된 torch를 사용하자.
4. torch와 torchvision은 서로 버전 호환성이 까다롭다. torch 버전에 맞는 torchvision을 설치할 것.
5. 또한 마찬가지로 pip를 이용하여 torchvision을 설치할 경우 작동하지 않는다. Build from source를 통해 Jetson내에서 새로 빌드할 것.
6. 이것은 Jetson 뿐만 아니라 공통적으로 적용되는 것인데, 어떤 작업이든 왠만하면 파이썬 가상환경 venv를 통해 작업을 하자. OS나 여러가지 프로그램에 의해 기본 파이썬 환경은 잘 작동하지 않는 경우가 많다. 이것 때문에도 정말 많은 오류가 난다.

자 그러면 설치를 시작해 보자.

1. Cuda 11.8 설치

먼저 나는 torch 2.1.0을 설치하고 싶기 때문에 CUDA 버전이 11.8이어야 한다. 현재 설치되어 있는 Jetpack의 버전을 살펴보자.

sudo apt-cache show nvidia-jetpack

Package: nvidia-jetpack
Version: 5.1.1-b56
....

5.1.1 버전이다. Jetpack 5.1.1 여기에 들어가보면 기본적으로 CUDA 11.4.19 가 설치되어 있다고 나온다. 그러므로 CUDA 11.8을 설치해야 한다.

아래와 같은 명령어를 입력하여 설치한다.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/arm64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-tegra-repo-ubuntu2004-11-8-local_11.8.0-1_arm64.deb
sudo dpkg -i cuda-tegra-repo-ubuntu2004-11-8-local_11.8.0-1_arm64.deb
sudo cp /var/cuda-tegra-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

자 그러면 이제 torch를 설치할 차례이다.

2. 파이썬 가상환경 설정

그 전에, 이전에 말했듯이, 가상환경에서 진행해야 한다. 안그러면 torch 패키지가 제대로 설치되지 않는다.

먼저 vritualenv를 설치해준다.

pip install virtualenv

그다음 가상환경을 만들어준다. yolo라는 이름으로 만들어보겠다.

python3.8 -m venv yolo

그 다음 가상환경을 활성화 시켜준다.

source yolo/bin/activate

이제부터 패키지를 설치하거나 실행할 때 모두 사전에 가상환경을 활성화 시켜주어야 한다.

3. torch 설치

내가 설치하고 싶은 것은 Jetson platform으로 빌드된 torch 2.1.0이다. 설치 파일은 여기 포럼에 올라와 있다.

포럼에 보면 아래와 같이 JetPack5로 빌드된 torch 2.1.0버전을 볼 수 있을것이다. 다운로드 받자.

다운로드가 완료되면 다운로드된 경로로 가서 pip로 설치한다. 가상환경 활성화 했는지 꼭 확인!

sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip3 install 'Cython<3'
pip install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

설치가 완료되었다면 CUDA가 잡히는지 확인해보자.

python
>>>import torch
>>>torch.cuda.is_available()
True

True가 나와야 한다. 안나온다면... 몰랑

4. torchvision 설치

torch가 설치되었다면 torchvision도 설치할 차례이다. torchvision은 github에서 직접 clone해서 build를 진행해야 한다. torch v2.1.0과 호환되는 torchvision은 v0.16.1이다.

아래와같은 순서로 설치를 진행한다. 가상환경 활성화 했는지 꼭 확인!

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch v0.16.1 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.16.0
python3 setup.py install

빌드하는데 시간이 꽤 걸릴것이다. 여기까지 완료했다면 아마 torchvision까지 설치가 완료되었을 것이다.

항상 느끼는거지만 개발 시작도 하기 전에 너무 힘들다 ㅜㅜ

Jetson tegra torch torchvision