Jetson보드에 Pytorch 설치하기

이전 포스트에서 Jetson 플랫폼에서 tensorflow설치하는 방법을 알아보았다

하지만 CUDA가 정상적으로 빌드되지 않은 것 같다

이번 포스트에서는 tensorflow와 호한되는 CUDA를 설치하려고 한다

필자가 설치한 tensorflow 버전은 2.12.0이기 때문에 이와 호환가능한 CUDA버전을 설치해야한다

 

 

jtop 명령어로 확인해보면 CUDA는 버전을 수정해야할 필요과 확실하다 (11.4 -> 11.8)

 

1. CUDA설치하기(Option)

참고 레퍼런스

https://developer.nvidia.com/cuda-11-8-0-download-archive?target_os=Linux&target_arch=aarch64-jetson&Compilation=Native&Distribution=Ubuntu&target_version=20.04&target_type=deb_local

CUDA Toolkit 11.8 Downloads

 

#Please ensure your device is configured per the CUDA Tegra Setup Documentation.

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

 

설치 결과 확인

jtop

 

정상적으로 설치가 완료 됬다

그리고 cuda가 잘 잡히는지 테스트 해보자

import torch

print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
print('cuDNN version: ' + str(torch.backends.cudnn.version()))

 

여전히 안잡힌다 Torch를 다시 설치해보자

 

2. Pytorch 설치하기 (핵심내용)

참고 레퍼런스 : https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html

위 링크 참고해서 본인이 맞는 jetpack 버전과 .whl 파일을 지정해주면 된다!

 

아래 예시 ) 본인은 jetpack 버전 5.1.1 사용

//torch 2.2.0
export TORCH_INSTALL=https://developer.download.nvidia.cn/compute/redist/jp/v511/pytorch/torch-2.0.0+nv23.05-cp38-cp38-linux_aarch64.whl
pip3 install --no-cache $TORCH_INSTALL


// torch 1.14
pip3 install --no-cache https://developer.download.nvidia.com/compute/redist/jp/v51/pytorch/torch-1.14.0a0+44dac51c.nv23.01-cp38-cp38-linux_aarch64.whl

 

 

테스트 결과 

정상적으로 설치가 완료된 것을 알 수 있다!

 

<Reference>

https://developer.nvidia.com/blog/simplifying-cuda-upgrades-for-nvidia-jetson-users/

https://www.tensorflow.org/install/source#tested_build_configurations

https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html