Install PyTorch for Nvidia Jetson Nano

[ deep-learning  pytorch  jetson  nano  ]

This is the step required for installing PyTorch on Nvidia Jetson Nano.

Select Python 3 as default Python.

sudo ln -s /usr/bin/python3 /usr/bin/python

Install Python 1.1.0

wget https://nvidia.box.com/shared/static/j2dn48btaxosqp0zremqqm8pjelriyvs.whl -O torch-1.1.0-cp36-cp36m-linux_aarch64.whl
sudo pip3 install numpy torch-1.1.0-cp36-cp36m-linux_aarch64.whl

Or Python 1.0.0

wget https://nvidia.box.com/shared/static/2ls48wc6h0kp1e58fjk21zast96lpt70.whl -O torch-1.0.0a0+bb15580-cp36-cp36m-linux_aarch64.whl
sudo pip3 install numpy torch-1.0.0a0+bb15580-cp36-cp36m-linux_aarch64.whl

Verify your installation

import torch
print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
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))

[Optional] Install torch2trt to enable PyTorch using TensorRT.

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
sudo python setup.py install

Why torch2trt

Please refer the table for the performance gap (FPS) for with/out TensorRT.

The results below show the throughput in FPS. You can find the raw output, which includes latency, in the benchmarks folder.

Model Nano (PyTorch) Nano (TensorRT) Xavier (PyTorch) Xavier (TensorRT)
alexnet 46.4 69.9 250 580
squeezenet1_0 44 137 130 890
squeezenet1_1 76.6 248 132 1390
resnet18 29.4 90.2 140 712
resnet34 15.5 50.7 79.2 393
resnet50 12.4 34.2 55.5 312
resnet101 7.18 19.9 28.5 170
resnet152 4.96 14.1 18.9 121
densenet121 11.5 41.9 23.0 168
densenet169 8.25 33.2 16.3 118
densenet201 6.84 25.4 13.3 90.9
densenet161 4.71 15.6 17.2 82.4
vgg11 8.9 18.3 85.2 201
vgg13 6.53 14.7 71.9 166
vgg16 5.09 11.9 61.7 139
vgg19     54.1 121
vgg11_bn 8.74 18.4 81.8 201
vgg13_bn 6.31 14.8 68.0 166
vgg16_bn 4.96 12.0 58.5 140
vgg19_bn     51.4 121

Reference

  • https://devtalk.nvidia.com/default/topic/1049071/jetson-nano/pytorch-for-jetson-nano/
  • https://github.com/NVIDIA-AI-IOT/torch2trt
Written on June 25, 2019