Jetson 系列开发板(NX/AGX /Nano)搭建pytorch-gpu环境

提醒:Jetson Xavier NX 用不了 nvidia-smi 命令

0.查看JetPack版本信息

sudo apt-cache show nvidia-jetpack

Package: nvidia-jetpack
Version: 4.6-b199
Architecture: arm64
Maintainer: NVIDIA Corporation
Installed-Size: 194

• 后面选择安装版本的时候需要根据JetPack版本信息选择版本

1.安装miniconda

• 下载地址:https://docs.conda.io/en/latest/miniconda.html

wget https://github.com/Archiconda/build-tools/releases/download/0.2.3/Archiconda3-0.2.3-Linux-aarch64.sh
bash Archiconda3-0.2.3-Linux-aarch64.sh

• 然后创建自己的虚拟环境即可

conda create -n base-jupiter python=3.6

2.安装pytorch-gpu

• 官网的下载地址并不包含aarch64适用的pytorch-gpu
• 以LTS (1.8.2)为例，官下载地址为：https://download.pytorch.org/whl/lts/1.8/torch_lts.html
• 打开可以验证发现并并不包含aarch64适用的pytorch-gpu
• 因此安装gpu版本的需要从NVIDIA官方进行下载，下载地址为：https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-11-now-available/72048

wget https://nvidia.box.com/shared/static/fjtbno0vpo676a25cgvuqc1wty0fkkg6.whl -O torch-1.10.0-cp36-cp36m-linux_aarch64.whl
pip install numpy Cpython
pip install torch-1.10.0-cp36-cp36m-linux_aarch64.whl

遇到问题及解决办法

• ImportError: libopenblas.so.0: cannot open shared object file: No such file or directory

sudo apt-get install libopenblas-dev

• OSError: libmpi_cxx.so.20: cannot open shared object file: No such file or directory

sudo apt-get install libopenmpi-dev 

• import torch 出现 Illegal instruction (core dumped)

vim ~/.bashrc

# 把以下内容加入到末尾扩充环境变量
export OPENBLAS_CORETYPE=ARMV8

source ~/.bashrc

3.安装torchvision

• 下载地址：https://github.com/pytorch/vision
• 版本对应关系

git clone -b v0.11.1 https://github.com/pytorch/vision.git vision-0.11.1
cd vision-0.11.1
export BUILD_VERSION=0.11.1
python setup.py install

4.效果测试

(base-jupiter) nvidia@nx:~$ python
Python 3.6.15 | packaged by conda-forge | (default, Dec  3 2021, 19:12:04)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.cuda.is_available()
True

参考资料

1. 查看Jetson系列产品JetPack的版本信息
2. NVIDIA Jetson Xavier NX搭建pytorch gpu环境（超详细）
3. NVIDIA JETSONTX2 安装 pytorch 出现错误：import torch 出现 Illegal instruction(core dumped)
4. ImportError: libopenblas.so.0: cannot open shared object file: No such file or directory
5. Jetson AGX Xavier安装Archiconda虚拟环境管理器与在虚拟环境中调用opencv
6. Jetson AGX Xavier安装torch、torchvision且成功运行yolov5算法
7. https://github.com/pytorch/vision