Jetson Xavier NX arm64 pytorch1.8.0安装

板子：NVIDIA Jetson Xavier NX

系统：Ubuntu 18.04.6 LTS (GNU/Linux 4.9.201-tegra aarch64)

$ uname -a
Linux nvidia-desktop 4.9.201-tegra #1 SMP PREEMPT Fri Feb 19 08:42:04 PST 2021 aarch64 aarch64 aarch64 GNU/Linux

## 显卡信息
$ lspci | grep -i nvidia
0004:00:00.0 PCI bridge: NVIDIA Corporation Device 1ad1 (rev a1)
0005:00:00.0 PCI bridge: NVIDIA Corporation Device 1ad0 (rev a1)

## cuda版本
$ cat /usr/local/cuda/version.txt 
CUDA Version 10.2.89

## tegra信息
$ cat /etc/nv_tegra_release
R32 (release), REVISION: 5.2, GCID: 27767740, BOARD: t186ref, EABI: aarch64, DATE: Fri Jul  9 16:05:07 UTC 2021

Jetson查看GPU使用jtop替代nvidia-smi， jtop输出：

按照官网的安装过程，选好配置后如下：

执行pip3 install torch torchvision torchaudio后测试安装结果：

$ python3
>>> import torch
>>> print(torch.__version__)
1.10.2
>>> print(torch.cuda.is_available())
False

可见cuda.is_available为false，这样无法使用GPU，查阅资料后可能是由于版本不对应导致的，版本对应关系可参考官网或者这个博主整理的CUDA驱动和CUDA Toolkit对应版本 。

在网上找来一圈博客后试了很多方法都失败，当我翻到一篇博客后，上面说到“pytorch官方并未直接提供arm平台的安装包”（这句话我没去考证是否确实如此），我才意识到我可能找错了方向，因为我用的板子是NVIDIA Jetson Xavier NX，跟直接在系统上装pytorch有很大的差别，于是重整旗鼓改变了google的搜索方向。

我在nvidia官网找到了归宿PyTorch for Jetson - version 1.10 now available

# ps:我这是root用户，非root用户遇到权限问题自行加sudo
## 使用conda切换一个新环境：
$ conda create -n python3.6 python=3
$ conda install python=3.6 # 安装python3.6
$ python3 -V
Python 3.6.9
$ pip3 install --upgrade pip # 更新pip
$ pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple/ # 添加清华大学源
$ mkdir tmp
$ cd tmp
# 我用的torch是1.8.0
$ wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
$ apt-get install libopenblas-base libopenmpi-dev
$ pip3 install Cython
$ pip3 install numpy torch-1.8.0-cp36-cp36m-linux_aarch64.whl
# 接着安装torchvision
$ apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
# 我用的torchvision是0.9.0
$ git clone --branch v0.9.0 https://github.com/pytorch/vision torchvision
$ cd torchvision
$ export BUILD_VERSION=0.9.0
$ python3 setup.py install --user

安装完后进行如下测试：

$ python3
>>> import torch
>>> print(torch.__version__)
1.8.0
>>> print('CUDA available: ' + str(torch.cuda.is_available()))
CUDA available: True
>>> print('cuDNN version: ' + str(torch.backends.cudnn.version()))
cuDNN version: 8000
>>> a = torch.cuda.FloatTensor(2).zero_()
>>> print('Tensor a = ' + str(a))
Tensor a = tensor([0., 0.], device='cuda:0')
>>> b = torch.randn(2).cuda()
>>> print('Tensor b = ' + str(b))
Tensor b = tensor([-0.2410,  0.9421], device='cuda:0')
>>> c = a + b
>>> print('Tensor c = ' + str(c))
Tensor c = tensor([-0.2410,  0.9421], device='cuda:0')
>>> import torchvision
>>> print(torchvision.__version__)
0.9.0

可以看到cuda.is_available为true了，此时就可以愉快的使用GPU了。

如在pyhton的import时遇到“Illegal instruction (core dumped)”报错，查看numpy版本，可能是1.19.5引发的问题，执行pip3 install numpy==1.19.4降级到1.19.4即可解决。

安装完毕！

所有的参考资料皆在原文中给了超链，此处不再一一归纳了。

接下来我将用GitHub上的yolov5导出GPU算法模型，博客写好后将再此处贴出链接。