Setting up PyTorch on NVIDIA Jetson boards

Today I was playing around with an NVIDIA Jetson Xavier NX when I came to realise that installing PyTorch on it is not as simple as doing:

pip install torch torchvision

As an acknowledgement I would like to thank https://cognitivexr.at/blog/2021/03/11/installing-pytorch-and-yolov5-on-an-nvidia-jetson-xavier-nx.html for getting it right but that article has some additional stuff, and some that is absent and I will try to keep it super straight-forward.

Please be wary that I have omitted some steps like cd to directory, creating a virtual environment, and other trivial things.

Installing PyTorch

PyTorch for L4T (the kernel patch for Jetson boards) is available separately and the .whl file has to be downloaded and installed via pip.

I used the link but you don’t need to go to it as I will summarise below.

I will install PyTorch 1.9.0, and you can visit the link for a different version, but the essential commands are still the same:

That will install torch.

To verify I imported torch in a python shell:

$ python
>> import torch
Illegal instruction (core dumped)

But I still had this error:

Illegal instructions (core dumped)

Thanks to the article I mentioned above, a simple fix exists. In your .bashrc, just add:

export OPENBLAS_CORETYPE=ARMV8

That fixes the issue.

$ python
>> import torch
>>

Now, to install torchvision, it seems like the only way to go about it is compiling it on the Jetson itself. I could be wrong, but it was quite easy to do.

There are some additional libraries that need to be installed. I install torchvision v0.10.0.

You can once again, verify by importing torchvision into Python shell.

$ python
>> import torch
>> import torchvision
>>

You can now run your ML models on the Jetson using PyTorch!

On the other hand, if you are looking to deploy ML models on the cloud, check out something we have built at our startup: https://deploif.ai. It lets you deploy models easily on AWS, Azure or GCP. It is as simple as drag-and-drop.

Thanks for reading!