Nvidia Jetson Xavier NX安装GPU版pytorch与torchvision

前提是已经安装好了系统，并通过JetPack配置完了cuda、cudnn、conda等库。

1. 安装GPU版pytorch

在base环境上新建环境，python版本3.8,激活并进入。

conda create -n pytorch_gpu python=3.8
conda activate pytorch_gpu

前往Nvidia论坛，下载Jetson NX专用的pytorch安装包。

传送门：https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048

我这里下载PyTorch v2.1.0

# 安装工具链
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip install Cython==0.29.21
pip install numpy
# 安装torch
pip install /path/to/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

2. 安装torchvision

Jetson上的gpu版本torchvision没有现成的安装包，需要自行编译安装。首先下载对应版本的torchvision

PyTorch v1.8 - torchvision v0.9.0
PyTorch v1.9 - torchvision v0.10.0
PyTorch v1.10 - torchvision v0.11.1
PyTorch v1.11 - torchvision v0.12.0
PyTorch v1.12 - torchvision v0.13.0
PyTorch v1.13 - torchvision v0.13.0
PyTorch v1.14 - torchvision v0.14.1
PyTorch v2.0 - torchvision v0.15.1
PyTorch v2.1 - torchvision v0.16.1
PyTorch v2.2 - torchvision v0.17.1
PyTorch v2.3 - torchvision v0.18.0

我这里下载torchvision v0.16.1

# 安装工具链
sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
# 下载torchvision
git clone --branch v0.16.1 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.16.1  # v0.16.1 为torchvision版本
python3 setup.py install --user

编译安装的过程较慢，静候编译安装结束。

3. 验证安装结果

新建一个python文件envtest.py，运行，验证安装结果。

# envtest.py
import torch
print('CUDA available: ' + str(torch.cuda.is_available()))
print('cuDNN version: ' + str(torch.backends.cudnn.version()))
a = torch.tensor([0., 0.], dtype=torch.float32, device='cuda')
print('Tensor a =', a)
b = torch.randn(2, device='cuda')
print('Tensor b =', b)
c = a + b
print('Tensor c =', c)

import torchvision
print(torchvision.__version__)

正常输出应如下：

CUDA available: True
cuDNN version: 8600
Tensor a = tensor([0., 0.], device='cuda:0')
Tensor b = tensor([ 0.4206, -1.0542], device='cuda:0')
Tensor c = tensor([ 0.4206, -1.0542], device='cuda:0')
0.16.1