Nvidia jetson系列 agx xavier/orin/nano/tx2 安装pytorch Gpu版

        记录一下工作中的问题。

        手头的 jetson  需要搭建深度学习环境，刚开始以为安装方式和正常Ubuntu环境一样，然而这个系统被改了很多，常规的安装方式和whl文件只能用pytorch cpu版本的，cuda用不了，想用cuda还是得装gpu版本，agx  pytorch_gpu只能通过英伟达官方发布的包安装，以下就记录一下详细安装过程。

        各jetson设备和版本安装方法类似，下面以pytorch1.9.0为例。

        博客提供whl文件，jetpack5.0以下：下载torch1.9.0或torch1.10.0；jetpack5.0：下载torch1.11.0或torch1.12 .0。

一、agx内置环境

1、xavier用官方sdkmanager刷机，包含各sdk。

2、jetpack4.6.1

3、cuda10.2

4、python3.6

二、torch安装

1、虚拟环境搭建

        如果不需要虚拟环境，直接看下一步。

        搭建虚拟环境可选择的比较多，比如anaconda（jetpack5.0可用），miniconda，archiconda等，前两个不建议使用，因为创建python3.6环境报错，这里推荐使用archiconda，

下载链接 https://github.com/Archiconda/build- tools/releases/download/0.2.3/Archiconda3-0.2.3-Linux-aarch64.sh

下载完成后执行

sh Archiconda3-0.2.3-Linux-aarch64.sh

按提示输入yes完成安装。

然后创建虚拟环境，选择python=3.6（官网提供的torch只有cp36的)

jetpack5.0  python=3.8（注意）

conda create -n pytorch python=3.6  #创建虚拟环境

conda activate pytorch  # 激活虚拟环境，此后torch和torchvision安装在虚拟环境中

2、文件下载

        可通过Jetson Zoo - eLinux.org网站自行下载，不过可能需要梯子翻墙。

jetpack5.0以下whl文件：

jetpack5.0以下whl文件-CSDN下载

torch-1.10.0a0+torchvision-0.11.0-cp36-cp36m-linux_aarch64

torch-1.9.0a0+torchvision-0.10.0-cp36-cp36m-linux_aarch64

torch-1.8.0a0+torchvision-0.9.0-cp36-cp36m-linux_aarch64

torch-1.7.0a0+torchvision-0.8.0-cp36-cp36m-linux_aarch64

jetpack5.0及以上版本whl文件：

jetpack5.0相关whl文件-CSDN下载

torch-1.13.0a0+torchvision-0.14.0-cp38-cp38m-linux_aarch64

torch-1.12.0a0+torchvision-0.13.0-cp38-cp38m-linux_aarch64

torch-1.11.0a0+torchvision-0.12.0-cp38-cp38m-linux_aarch64

torch-1.14.0a0+44dac51c.nv23.01-cp38-cp38-linux_aarch64.whl

3、安装

sudo apt-get install python3-pip libopenblas-base libopenmpi-dev

sudo pip3 install Cython 

cd /下载好文件的文件夹

pip3 install torch-1.9.0-cp36-cp36m-linux_aarch64.whl

pip3 install numpy

测试一下，应该会报错，如下。

不过也很简单，修改一下环境变量就可以了。

sudo gedit ~/.bashrc   #编辑bashrc

export OPENBLAS_CORETYPE=ARMV8  #加到最后一行

source ~/.bashrc   #更新一下环境变量

再测试一下

torch中cuda能用了，完美解决。 

三、安装torchvision

        torch1.9.0 对应torchvision0.10.0 ，torchvision没有官方whl文件，不过有两种方式可以安装，第一种比较简单，直接安装非官方编译好的whl文件，比较简单；第二种需要编译安装，比较麻烦。

1、torchvision安装方法一

pip3 install torchvision-0.10.0 -cp36-cp36m-linux_aarch64.whl

测试一下，torchvision版本号不知道为啥怪怪的。

2、torchvison安装方法二

        以下是官方给出安装方式。

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev

git clone --branch release/0.10 https://github.com/pytorch/vision torchvision 

cd torchvision export BUILD_VERSION=0.10.0

sudo python3 setup.py install --user

这个过程较长，耐心等待，然后安装完成，你以为真的完了，想多了，哈哈哈，大概率会报错，如下图。

不过不要慌，找到问题就能解决，也很简单，更新一下pillow就ok了。

 pip3 install --upgrade pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

四、虚拟环境导入opencv

虽然jetson自带opencv，但是在虚拟环境中无法使用，需要配置一下。

如果需要opencv，找到 /usr/lib/python3.6/dist-packages/cv2/python-3.6路径下的cv2.cpython-36m-aarch64-linux-gnu.so文件，复制到路径/home/nvidia/archiconda3/envs/pytorch/lib/python3.6文件夹下就可以了。

五、虚拟环境导入tensorrt

和opencv情况相同。

如果需要tensorrt，找到 /usr/lib/python3.6/dist-packages/tensorrt路径下的tensorrt.so文件，复制到路径/home/nvidia/archiconda3/envs/pytorch/lib/python3.6文件夹下就可以了。
 

以上就是在英伟达 agx 上安装pytorch_gpu版本安装过程，踩坑无数，花了很长时间。

欢迎交流指正，私信我就行了。