系统信息:
板子: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算法模型,博客写好后将再此处贴出链接。