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

NVIDIA Jetson Xavier NX搭建pytorch gpu环境(超详细)_第1张图片

注意选择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  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

为对应的torchvision版本

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

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

NVIDIA Jetson Xavier NX搭建pytorch gpu环境(超详细)_第2张图片

验证torch gpu版本安装完成

python
import torch
print(torch.cuda.is_available())

返回True即可

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

你可能感兴趣的:(anaconda,python,pytorch)