nvidia jetson设备（Jetson Nano, TX1/TX2,Xavier NX/AGX, AGX Orin ）以及树莓派如何安装cpu/cuda版的pytorch和torchvision

0. 前言

以下安装以Jetson Nano为例

如果你的nano是python3.6.9（对应jetpack 4.6），千万不要尝试升级python版本。nano最高只支持jetpack4.6。

其余的拓展到其他的jetson设备以及其他的微型电脑如树莓派

1. 安装步骤

环境：jetson nano ubuntu18.04，已经刷好了jetpack4.6，python3是3.6.9
想安装cpu版的torch很简单，直接输pytorch上的官网命令就行。
但如果想用cuda，也就是gpu版本的torch，只有两种方法：

1. 自己编译pytorch的源码。
2. 下载英伟达官方或者别人预编译好的pytorch的轮子。

这边我们只说第二种方法：

2. 面向nano和树莓派

这里参照以下开源网站进行安装即可。包含torch和torchvision
https://github.com/Qengineering/PyTorch-Jetson-Nano

在此处按步骤进行安装
https://qengineering.eu/install-pytorch-on-jetson-nano.html

3. 面向所有jetson设备

（包含torch，不包含torchvision）
进入英伟达网站：https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048

3.1 参照

按照jetpack版本以及python版本进行安装即可

3.2 按版本安装torch

这里需要特别注意的一个问题就是你需要选择与你机器上jetpack相匹配的版本，千万别下错了，比如我这里用的是jetpack4.6，那我安装的pytorch就是V1.10.0版本。

下载完pytorch的whl文件后，首先apt-get install一些环境必须的前置程序，然后进入到你创建的conda环境下，使用pip命令安装Cython和pytorch。

sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
pip install Cython
pip install numpy xxx.whl

3.3 验证是否安装成功

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))

若上述程序执行无误，证明pytorch安装成功，并且可以调用jetson的GPU来进行运算。

3.3 安装torchvision

这里要注意torchvision的版本和pytorch是严格对应的。

比如我们的1.10 对应的就是v0.11.1，那么在拉取代码仓和install环节都需要制定版本。以下命令包含了代码拉取及安装，install过程时间很长，耐心等待

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch v0.11.1 https://github.com/pytorch/vision torchvision   # see below for version of torchvision to download
cd torchvision
export BUILD_VERSION=0.11.1  # where 0.x.0 is the torchvision version  
python3 setup.py install --user

验证torchvision

import torchvision
print(torchvision.__version__)

4. 注意事项

版本一定要对应。以nano为例，下载的时候注意，jetson nano只能下3.6python支持的，也就是最高支持pytorch的1.10.0版本，对应torchvision的版本是0.11.0。为什么呢？是因为

Hi @m.ahmad4984, JetPack 5.0.2 is in the final preparations to be
released, so it should be posted in the next few days pending any
unforeseen delays or blockers. As per the Jetson Software Roadmap 52,
JetPack 4.6 was the last major release for Nano/TX1/TX2 and JetPack
5.x supports Xavier and Orin devices.

官方原话，意思是nano最高支持的jetpack版本就是4.6，也就是python3.6.9，torch最后一个支持3.6的版本就是1.10.0。如果你不是nano是xavier或者更高级的型号，那就选python3.8的版本。

5. 常见问题

1. 下载完轮子后用pip3安装，如果提示：

“error：torch is not a supported wheel on this platform”

如果你下载的是预编译好的，出现了这个问题，实际上是你的pip3版本和你的python3.6不兼容，我这边之前尝试更新python和pip的时候遗留下来了这种问题。

解决办法

wget https://bootstrap.pypa.io/get-pip.py
sudo python3.6 get-pip.py

即重新安装对应版本的pip。

2. 进入python3环境，尝试import torch时出现：

numpy.core.multiarray failed to import
或者
numpy illegal instruction (core dumped) jetson

意思就是你的numpy版本有问题，和torch不兼容！

解决方法：
首先卸载所有的numpy，再重新安装低版本numpy：

#1.19.5出错
pip3 install numpy==1.19.4

3. git clone安装torchvision时候出现网络错误

Fatal: Unable To Access ‘Https://Github.Com/…‘

解决办法：

要么是修复这个git的网络问题或者等待（有时会恢复正常）

要么使用轮子安装：
https://github.com/KumaTea/pytorch-aarch64/releases
电脑中下载后再放到jetson上

6. 安装完毕的检验

输入以下命令查看：

import numpy
import torch
torch.__version__
torch.cuda.is_available()

如果没报异常且torch.cuda.is_available显示True，说明安装正常完成了，如下图所示：