torch / vision 설치 on Jetson Xavier (yolo GPU학습용)

 

1. 기본 설치

$ sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
$ sudo pip3 install 'Cython<3'
$ sudo pip3 install numpy

 

2. torch 다운로드

   jetpack 5.1.3이 설치된 Xavier에서 여러가지 방법을 해봤다. 소스로 설치를 해봤고, whl파일로 설치를 해보기도 했고... xavier를 완전히 초기화한 상태에서 다시 작업하면서, whl 파일을 다운로드 받아서 설치해보는 방법으로 진행하고 있다.

  다운로드는 아래의 주소에서 받을 수 있다.

https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048

PyTorch for Jetson

 

jetpack5.1.3이 설치된 것이니, PyTorch v2.1.0 을 다운로드 받았다. 안내된 내용에는 jetpack 5.1.2까지만 언급되어있다.

• JetPack 5.1 (L4T R35.2.1) / JetPack 5.1.1 (L4T R35.3.1) / JetPack 5.1.2 (L4T R35.4.1)
  • Python 3.8 - torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

 

3. pytorch 설치

$ sudo pip3 install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl 

Processing ./torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
Collecting filelock
  Downloading filelock-3.15.4-py3-none-any.whl (16 kB)
Collecting sympy
  Downloading sympy-1.12.1-py3-none-any.whl (5.7 MB)
     |████████████████████████████████| 5.7 MB 3.2 MB/s 
Collecting networkx
  Downloading networkx-3.1-py3-none-any.whl (2.1 MB)
     |████████████████████████████████| 2.1 MB 3.3 MB/s 
Collecting jinja2
  Downloading jinja2-3.1.4-py3-none-any.whl (133 kB)
     |████████████████████████████████| 133 kB 3.1 MB/s 
Collecting fsspec
  Downloading fsspec-2024.6.1-py3-none-any.whl (177 kB)
     |████████████████████████████████| 177 kB 3.6 MB/s 
Collecting typing-extensions
  Downloading typing_extensions-4.12.2-py3-none-any.whl (37 kB)
Collecting mpmath<1.4.0,>=1.1.0
  Downloading mpmath-1.3.0-py3-none-any.whl (536 kB)
     |████████████████████████████████| 536 kB 3.0 MB/s 
Collecting MarkupSafe>=2.0
  Downloading MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl (26 kB)
Installing collected packages: filelock, mpmath, sympy, networkx, MarkupSafe, jinja2, fsspec, typing-extensions, torch

  Attempting uninstall: MarkupSafe
    Found existing installation: MarkupSafe 1.1.0
    Not uninstalling markupsafe at /usr/lib/python3/dist-packages, outside environment /usr
    Can't uninstall 'MarkupSafe'. No files were found to uninstall.
Successfully installed MarkupSafe-2.1.5 filelock-3.15.4 fsspec-2024.6.1 jinja2-3.1.4 mpmath-1.3.0 networkx-3.1 sympy-1.12.1 torch-2.1.0a0+41361538.nv23.6 typing-extensions-4.12.2

 

4. pytorch 설치 확인

$ python3

Python 3.8.10 (default, Nov 22 2023, 10:22:35) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True
>>> quit()

True가 출력되면 정상 설치된 것이다.

 

5. torchvision 설치전 패키지 설치

$ sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev

 

6. torchvision 다운로드

위에서 설치한 pytorch 버전에 따라서, 설치할 torchvision의 버전 짝이 있다.  위에서는 pyTorch v2.1을 설치한 것이라, torchvision v0.16.1을 다운로드 받아 설치해야 한다.

PyTorch v1.12 - torchvision v0.13.0
PyTorch v1.13 - torchvision v0.13.0
PyTorch v1.14 - torchvision v0.14.1
PyTorch v2.0 - torchvision v0.15.1
PyTorch v2.1 - torchvision v0.16.1
PyTorch v2.2 - torchvision v0.17.1
PyTorch v2.3 - torchvision v0.18.0

git으로 다운로드 한다.

$ git clone --branch v0.16.1 https://github.com/pytorch/vision torchvision-0.16.1

 

7. torchvision 설치

설치 전에 환경변수를 설정한 후, 설치를 진행한다.   nvidia forum에 등록된 글에는 user 홈에 설치하도록 안내되어있다.($python3 setup.py install --user)

$ cd torchvision-0.16.1
$ export BUILD_VERSION=0.16.1
$ sudo python3 setup.py install

 

한참의 complie 후에 아래와 같이 설치 완료된다.

....

Using /usr/lib/python3/dist-packages
Searching for numpy==1.17.4
Best match: numpy 1.17.4
Adding numpy 1.17.4 to easy-install.pth file
Installing f2py script to /usr/local/bin
Installing f2py3 script to /usr/local/bin
Installing f2py3.8 script to /usr/local/bin

Using /usr/lib/python3/dist-packages
Searching for filelock==3.15.4
Best match: filelock 3.15.4
Adding filelock 3.15.4 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for sympy==1.12.1
Best match: sympy 1.12.1
Adding sympy 1.12.1 to easy-install.pth file
Installing isympy script to /usr/local/bin

Using /usr/local/lib/python3.8/dist-packages
Searching for networkx==3.1
Best match: networkx 3.1
Adding networkx 3.1 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for typing-extensions==4.12.2
Best match: typing-extensions 4.12.2
Adding typing-extensions 4.12.2 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for fsspec==2024.6.1
Best match: fsspec 2024.6.1
Adding fsspec 2024.6.1 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for jinja2==3.1.4
Best match: jinja2 3.1.4
Adding jinja2 3.1.4 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for mpmath==1.3.0
Best match: mpmath 1.3.0
Adding mpmath 1.3.0 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Searching for MarkupSafe==2.1.5
Best match: MarkupSafe 2.1.5
Adding MarkupSafe 2.1.5 to easy-install.pth file

Using /usr/local/lib/python3.8/dist-packages
Finished processing dependencies for torchvision==0.16.1+fdea156

 

8. torchvision 설치 확인

python프로그램을 만들어서 실행시켜 보자.

$ vi torch_test.py

import torch
import torchvision

def check_torch_and_torchvision():
    try:
        # Check if torch is installed
        print(f"PyTorch version: {torch.__version__}")
        
        # Check if torchvision is installed
        print(f"Torchvision version: {torchvision.__version__}")
        
        # Create a tensor
        tensor = torch.tensor([1.0, 2.0, 3.0, 4.0])
        print(f"Created tensor: {tensor}")
        
        # Perform a simple tensor operation
        tensor_sum = tensor.sum()
        print(f"Sum of tensor elements: {tensor_sum}")
        
        # Check for CUDA support
        if torch.cuda.is_available():
            print("CUDA is available. GPU is supported.")
            print(f"CUDA version: {torch.version.cuda}")
            print(f"Number of CUDA devices: {torch.cuda.device_count()}")
            print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
        else:
            print("CUDA is not available. GPU is not supported.")
        
        print("PyTorch and Torchvision are installed and working correctly.")
    except ImportError as e:
        print(f"Error: {e}")
        print("PyTorch or Torchvision is not installed or there is an issue with the installation.")

if __name__ == "__main__":
    check_torch_and_torchvision()

 

실행한 결과는 아래와 같다.

$ python3 python_test.py
PyTorch version: 2.1.0a0+41361538.nv23.06
Torchvision version: 0.16.1+fdea156
Created tensor: tensor([1., 2., 3., 4.])
Sum of tensor elements: 10.0
CUDA is available. GPU is supported.
CUDA version: 11.4
Number of CUDA devices: 1
CUDA device name: Xavier
PyTorch and Torchvision are installed and working correctly.