Jetson Xavier NX - Python virtual environment and ML platforms(tensorflow, Pytorch) installation

This article is in line with the Jetson Xavier NX-JetPack 4.4 (production release) headless setup described in the previous blog.

So far, I haven't used the Python virtual environment while using the Jetson Nano. The reason is that Jetson Nano only supports SD cards, so it is convenient to prepare several SD cards and install and use the necessary packages. And because JetPack is updated frequently, it was also a big reason to prepare a new image on an SD card from time to time. However, the Xavier NX does offer an M.2 slot for SSD installation. SSDs offer significantly better performance than SD cards or eMMC memory.  Therefore, I will install and use a SSD on the Xavier NX.  When using the Jetson Nano, replacing the SD card was like using a completely new system. However, if you installed the package or saved data using an SSD, it is convenient to continue using the SSD regardless of the SD card replacement.

I will be mainly working on AI edge computing on Xavier NX. And I will use Python 3 as the development language. In the future, most of the data or packages will be installed on SSDs to improve performance and to enable the system to be used without replacing SD cards.

Creating a Python virtual environment

You can install TensorFlow, PyTorch, Yolo, etc. frameworks for use in one virtual environment, or you can create multiple Python virtual environments to install each framework in a separate virtual environment. Just choose the method that suits your Hankyung. Sometimes I feel the need to use TensorFlow and PyTorch together, so I will install multiple machine learning frameworks on a single virtual machine.

spypiggy@XavierNX:~$ sudo apt-get install python3-dev python3-pip
spypiggy@XavierNX:~$ pip3 install  -U virtualenv

The advantage of the Python virtual environment is that all packages installed in the virtual environment are installed in the virtual environment, so if you delete the virtual environment, everything installed in the virtual environment is cleared and does not affect other Python virtual environments.

Make sure to create the Python virtual environment on the SSD installed earlier. Previously, the SSD was mounted at /home/spypiggy. Therefore, create a virtual environment under /home/spypiggy.

Create a Python virtual environment with the following command. When you run this command, the core elements of your current Python 3 are copied to the /home/spypiggy/python directory.

spypiggy@XavierNX:~$ python3 -m virtualenv --system-site-packages /home/spypiggy/python

To enter the virtual environment, use the source command. Notice that the prompt name has been changed.

spypiggy@XavierNX:~$ source /home/spypiggy/python/bin/activate
(python) spypiggy@XavierNX:~$ 

From now on, all packages installed using the pip3 command are installed under the /home/spypiggy/python directory, and these packages are effective only when working in a virtual environment.

Install machine learning related platforms

Tensorflow

Do not follow the installation method of Tensorflow homepage. NVidia provides TensorFlow optimized for Jetpack version.
You can download the appropriate version for JetPack4.4 from the https://developer.download.nvidia.com/compute/redist/jp/v44/tensorflow/ page.
If the 4.6 version will be released after a while, you can probably search https://developer.download.nvidia.com/compute/redist/jp/v46/tensorflow/.

The naming convention is as follows.

TensorFlow version + NVIDIA container version of TensorFlow + Python version + CPU.

It is safe to select the latest version with a higher NVIDIA container version of TensorFlow value from the same TensorFlow version.

Select the desired version from TensorFlow 2.X version and 1.X version, and then install it. If you want to use both, you can create a new virtual environment and install it in another virtual environment. I will install tensorflow-1.15.2+nv20.6-cp36-cp36m-linux_aarch64.whl.

(python) spypiggy@XavierNX:~$ sudo apt-get install -y libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev
(python) spypiggy@XavierNX:~$ pip3 install numpy grpcio absl-py py-cpuinfo psutil portpicker six mock requests gast h5py astor termcolor protobuf keras-applications keras-preprocessing wrapt google-pasta(python) spypiggy@XavierNX:~$ wget https://developer.download.nvidia.com/compute/redist/jp/v44/tensorflow/tensorflow-1.15.2+nv20.6-cp36-cp36m-linux_aarch64.whl
(python) spypiggy@XavierNX:~$ wget https://developer.download.nvidia.com/compute/redist/jp/v44/tensorflow/tensorflow-1.15.2+nv20.6-cp36-cp36m-linux_aarch64.whl
(python) spypiggy@XavierNX:~$ pip3 install tensorflow-1.15.2+nv20.6-cp36-cp36m-linux_aarch64.whl

After installation, you can check as follows.

(python) spypiggy@XavierNX:~$ python
Python 3.6.9 (default, Apr 18 2020, 01:56:04) 
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
2020-07-22 09:45:49.770178: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.2
WARNING:tensorflow:Deprecation warnings have been disabled. Set TF_ENABLE_DEPRECATION_WARNINGS=1 to re-enable them.
>>> tf.__version__
'1.15.2'

PyTorch

How to install PyTorch 1.6 on Jetson Nano using JetPack 4.4 (DP Version) is explained in this blog. How to install PyTorch on Xavier NX is not much different from the content of this blog.

Like TensorFlow, NVidia provides an installation package optimized for Jetson. Install this package.

We always use Python 3.X. Therefore, download the whl file that can be used in Python 3.6. And install the necessary packages as follows. JetPack 4.3 and later comes with OpenCV 4.1.1 and numpy. So you don't have to install numpy anymore.

(python) spypiggy@XavierNX:~$ sudo apt-get update
(python) spypiggy@XavierNX:~$ wget https://nvidia.box.com/shared/static/yr6sjswn25z7oankw8zy1roow9cy5ur1.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
(python) spypiggy@XavierNX:~$ sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
(python) spypiggy@XavierNX:~$ pip3 install Cython
(python) spypiggy@XavierNX:~$ pip3 install torch-1.6.0-cp36-cp36m-linux_aarch64.whl

Download Torchvision for Jetson Xavier NX

If you have successfully installed PyTorch 1.6.0, install Torchvision 0.7.0. The latest version of torchvision can be found at https://github.com/pytorch/vision/releases.

(python) spypiggy@XavierNX:~$ sudo apt-get install libjpeg-dev zlib1g-dev
(python) spypiggy@XavierNX:~$ wget https://github.com/pytorch/vision/archive/v0.6.1.tar.gz   
(python) spypiggy@XavierNX:~$ cd v0.7.0
(python) spypiggy@XavierNX:~/v0.7.0$ python3 setup.py install
(python) spypiggy@XavierNX:~/v0.7.0$ sudo apt-get install libfreetype6-dev
(python) spypiggy@XavierNX:~/v0.7.0$ pip3 uninstall pillow
(python) spypiggy@XavierNX:~/v0.7.0$ pip3 install --no-cache-dir pillow 

Let's check whether the installation is correct.
If you see the screen like this, the installation is successful.

root@jetpack-4:/usr/local/src/torchvision# python3
Python 3.6.9 (default, Apr 18 2020, 01:56:04)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torchvision
>>> torch.__version__
'1.6.0'
>>> torchvision.__version__
'0.7.0
>>>

Wrapping up

So far, in Xavier NX, I have created a Python virtual environment in SSD storage and installed the machine learning platform Tensorflow and PyTorch. Unlike Jetson Nano, Xavier NX can be equipped with SSD. It is recommended to use SSD actively to improve performance. To install as many Python packages as possible in the SSD repository, I created a Python virtual environment and installed TensorFlow, PyTorch, etc.

I'll run a lot of tests that I've done on Jetson Nano so far on Xavier and post articles comparing performance.