NVIDIA Jetson Xavier NX搭建pytorch gpu环境（超详细）

NVIDIA® Jetson Xavier™ NX开发套件在搭建tensorflow-gpu环境时可以使用指令直接安装或者官网下载whl文件安装。作者在安装pytorch环境时总是安装不上gpu版本

报错 AssertionError: Torch not compiled with CUDA enabled

由于该板卡内核为aarch64，没有找到相应的安装包，所以总是出错，在此记录避坑。

1、环境管理

可以使用miniconda3，当然若一张存储卡只是在一个场景中应用直接搭建环境即可（应该不会有什么影响，欢迎指正）。

miniconda3安装

annaconda官网找到对应版本，下载。链接 Miniconda — Conda documentation

注意选择aarch64对应版本。

将下载好的文件上传至板卡，作者使用WinSCP。

pip install "xxxxx"安装

 然后yes或者回车接受即可。

2、torch gpu环境搭建

在pytorch官网找不到GPU版本的离线包，只有CPU版本的，若要安装CPU版本的，直接在官网下载即可。链接 https://download.pytorch.org/whl/torch_stable.html

注意选择aarch64版本

若安装GPU版本，NVIDIA官方给出了安装提示。

下载链接 

PyTorch for Jetson - version 1.10 now available - Jetson Nano - NVIDIA Developer Forums

注意python、torch以及torchvision版本对应关系

 下载完成后pip3 install xxxx安装即可

然后安装对应版本torchvision

指令如下

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
$ git clone --branch <version> https://github.com/pytorch/vision torchvision   
$ cd torchvision
$ export BUILD_VERSION=0.x.0  # where 0.x.0 is the torchvision version  
$ python3 setup.py install --user

<version>为对应的torchvision版本

python3 setup.py install --user指令执行时间可能略长

torch、torchvision以及python版本对应关系表

验证torch gpu版本安装完成

python

import torch

print(torch.cuda.is_available())

返回True即可

若返回False说明没有安装好GPU版本pytorch。