Skip to content

README.md

Latest commit

 

History

History
133 lines (94 loc) · 7.56 KB

README.md

File metadata and controls

133 lines (94 loc) · 7.56 KB

Pytorch

Introduction

Deploying PyTorch models to embedded edge devices is a critical step in bringing AI applications to life. NVIDIA's Jetson platform, with its powerful GPU computing capabilities and comprehensive AI software stack, has become an ideal choice for running PyTorch models.

However, because Jetson is based on the ARM architecture, which differs from common x86 server environments, setting up a PyTorch environment on it cannot be accomplished with a simple pip install command. Developers often face challenges such as finding the correct version of pre-compiled packages, managing complex dependencies, and performing necessary performance optimizations.

This article aims to provide a clear and practical guide, focusing on how to quickly and correctly configure the PyTorch environment on the Jetson platform, helping you kickstart your PyTorch development journey on Jetson.


Image from： pypi

Installing PyTorch on reComputer Nvidia Jetson

Set Up Your Environment

  • JetPack 5/6: Make sure you have NVIDIA JetPack 5 or 6 installed on your reComputer. JetPack includes the necessary libraries and tools for developing on NVIDIA Jetson platforms.

  • CUDA: Verify that CUDA is installed and properly configured. PyTorch relies on CUDA for GPU acceleration. Ensure that the CUDA version installed is compatible with the PyTorch version you plan to install.

Type cat /etc/nv_tegra_release and nvcc -V in your terminal. If the returned content is similar to the screenshot below, it indicates that the corresponding environment has been properly installed in your Jetson.

Installing PyTorch Using a .whl File

To install PyTorch on your reComputer with the specified JetPack and CUDA versions, follow these steps:

Download the PyTorch Wheel File

Choose the correct wheel file based on your JetPack, CUDA and python version:

Install the Wheel File

  1. Open a Terminal:

    • Navigate to the directory where you downloaded the .whl file.

  2. Install:

    sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip3 install 'Cython<3'
pip3 install numpy
sudo pip3 install <filename>.whl

    Replace <filename> with the name of the downloaded .whl file.

Verify Installation

To verify that PyTorch has been installed correctly on your system, launch an interactive Python interpreter from the terminal and run the following commands:

```python
import torch
print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
print('cuDNN version: ' + str(torch.backends.cudnn.version()))
a = torch.cuda.FloatTensor(2).zero_()
print('Tensor a = ' + str(a))
b = torch.randn(2).cuda()
print('Tensor b = ' + str(b))    
c = a + b
print('Tensor c = ' + str(c))
```
```python
import torchvision
print(torchvision.__version__)
```

More Tutorial Content

Tutorial Type Description
Official PyTorch Tutorial doc An official PyTorch tutorial that provides a complete learning path.
PyTorch Development Documentation doc Official PyTorch development documentation provided by PyTorch.