Jetson Nano 安装torch和torchvision

最近在玩Jetson Nano，环境都部署好了，并且将教程上的 jetson-inference例子跑了一遍，想自己动手训练一下自己的模型。需要安装pytorch，记录一下安装过程

安装pytorch

安装pytorch可以使用官方提供的whl，官网上有介绍各个版本对应关系--可参考文末图片附件，根据自己的jetpack版本选择对应的pytorch就好，注意，需要科学上网
当然此次分享出我安装的pytorch包，我下载的是pytorch v1.6.0
链接：https://pan.baidu.com/s/1Mhf7CCpJbdKp0seng4Mu4Q
提取码：rf07

如果可以科学上网可以使用以下命令直接安装：

wget https://nvidia.box.com/shared/static/9eptse6jyly1ggt9axbja2yrmj6pbarc.whl -O torch-1.6.0-cp36-cp36m-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev 
pip3 install Cython
pip3 install numpy torch-1.6.0-cp36-cp36m-linux_aarch64.whl

安装torchvision

torchvision 安装要根据torch版本来选择
Select the version of torchvision to download depending on the version of PyTorch that you have installed

图片.png

$ sudo apt-get install libjpeg-dev zlib1g-dev
$ git clone --branch v0.7.0 https://github.com/pytorch/vision torchvision   # see below for version of torchvision to download
$ cd torchvision
$ export BUILD_VERSION=0.7.0  # where 0.x.0 is the torchvision version  
$ sudo python setup.py install     # use python3 if installing for Python 3.6
$ cd ../  # attempting to load torchvision from build dir will result in import error
$ pip install 'pillow<7' # always needed for Python 2.7, not needed torchvision v0.5.0+ with Python 3.6

验证

安装完成后，使用如下命令来验证

>>> 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))

>>> import torchvision
>>> print(torchvision.__version__)

补充

补充一下windows 安装pytorch的方法
我发现并不能直接使用pip install torch的方式来安装torch， 会报错；于是在google了一下后，找到了安装方法
打开pytorch官网， 找到如下内容

image.png

然后运行命令：

pip install torch==1.7.0+cpu torchvision==0.8.1+cpu torchaudio===0.7.0 -f https://download.pytorch.org/whl/torch_stable.html

验证一下：

image.png

附

附上torch版本和jetson nano中其他配置的版本对应关系

图片.png