Jetson上的pytorch国内源下载和torchvision安装教程

【安装前提】

jetson上系统刷机前必须把cuda刷进去，否则安装pytorch是用不了cuda的，且不能更换系统自带python版本也不能更换手动自带cuda否则pytorch都是不能正常使用的。

【安装步骤】

首先安装依赖：

复制代码

sudo apt install libopenblas-dev libatlas-base-dev

前往PyTorch for Jetson，下载所安装的jetpack版本支持的最高版本的torch wheel 安装包到Downloads目录下。注意一定要对应自己操作系统安装。比如我的jetson是jetpack5.1.x对应下图中红框的torch安装包，需注意Python 版本为 3.8。那么就找文件名包含cp38的，如果下载不下来可以尝试国内源下载gitee.com/FIRC/pytorch-jetson-cn-mirror或者github搜pytorch-jetson-cn-mirror，例如我找到如下whl对应我的jetpack5

torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

下载完后保存到本地，切换到对应目录

复制代码

pip install torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl

安装即可。然后进行验证

复制代码

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

如果返回True则安装完成，False就要检查自己cuda是否安装或者更换过。

接下来安装torchvision,首先安装依赖；

复制代码

sudo pip3 install numpy requests Pillow
sudo apt install libjpeg-dev libpng-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev

由于torhchvision早期版本都是没有提供whl文件因此需要从源码编译，torch2.3.0已经开始提供torchvision了，如果你用jetpack6可能就不需要从源码开始安装。这里以源码安装为例子

首先查询torch对应torchvision版本

下载源码：

复制代码

git clone --branch v0.16.2 https://github.com/pytorch/vision

网络不行clone慢的话，直接下载压缩包到PC

再上传jetson，解压即可

复制代码

unzip vision-0.16.2.zip

编译安装torchvision

复制代码

cd vision-0.16.2                    # 进入torchvision目录
export BUILD_VERSION=0.16.2         # 将BUILD_VERSION环境变量设置为值 0.16.2  
sudo python3 setup.py install

需要等待30min左右，出现以下提示则安装成功

安装成功后退出torchvision的安装目录再import torchvision进行验证，否则会出现以下warning

复制代码

(pytorch) nx@nx-desktop:~/Downloads/vision-0.15.2$ python
Python 3.8.18 (default, Sep 11 2023, 13:19:25) 
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torchvision
/home/nx/Downloads/vision-0.15.2/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: ''If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?
  warn(
/home/nx/Downloads/vision-0.15.2/torchvision/__init__.py:25: UserWarning: You are importing torchvision within its own root folder (/home/nx/Downloads/vision-0.15.2). This is not expected to work and may give errors. Please exit the torchvision project source and relaunch your python interpreter.
  warnings.warn(message.format(os.getcwd()))