Install PyTorch for Nvidia Jetson Nano
[jetson
deep-learning
pytorch
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