玩转NVIDIA Jetson （25）— jetson 安装pytorch和torchvision

在jetson上安装pytorch能够快速的帮助PC用户在熟悉的环境下进行开发，但是这个过程里边有一些坑，我建议大家参照官方给出的安装指导来进行安装，包括pytorch和包含各种网络的torchvision。官方安装如下，这个topic里会持续更新各个版本的pytorch。

PyTorch for Jetson – version 1.11 now available – Jetson & Embedded Systems / Jetson Nano – NVIDIA Developer Forums

1.安装miniforge虚拟环境

我还是建议大家用虚拟环境，不然jetson上的python环境非常容易被搞乱，具体安装方法见之前的文章。

2.安装pytorch

这里需要特别注意的一个问题就是你需要选择与你机器上jetpack相匹配的版本，千万别下错了，比如我这里用的是jetpack4.6，那我安装的pytorch就是V1.10.0版本。

下载完pytorch的whl文件后，首先apt-get install一些环境必须的前置程序，然后进入到你创建的conda环境下，使用pip命令安装Cython和pytorch。

sudo apt-get install libopenblas-base libopenmpi-dev libomp-dev
pip install Cython
pip install numpy xxx.whl

 验证是否安装成功

import torch
print(torch.__version__)
print('CUDA available: ' + str(torch.cuda.is_available()))
print('cuDNN version: ' + str(torch.backends.cudnn.version()))
a = torch.cuda.FloatTensor(2).zero_()
print('Tensor a = ' + str(a))
b = torch.randn(2).cuda()
print('Tensor b = ' + str(b))
c = a + b
print('Tensor c = ' + str(c))

若上述程序执行无误，证明pytorch安装成功，并且可以调用jetson的GPU来进行运算。

3.安装torchvision

这里要注意torchvision的版本和pytorch是严格对应的。

比如我们的1.10 对应的就是v0.11.1，那么在拉取代码仓和install环节都需要制定版本。以下命令包含了代码拉取及安装，install过程时间很长，耐心等待

$ sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
$ git clone --branch v0.11.1 https://github.com/pytorch/vision torchvision   # see below for version of torchvision to download
$ cd torchvision
$ export BUILD_VERSION=0.11.1  # where 0.x.0 is the torchvision version  
$ python3 setup.py install --user

验证torchvision

import torchvision
print(torchvision.__version__)

这个时候就可以尽情使用pytorch自带的众多训练好的神经网络做迁移学习了。