NVIDIA ORIN配置Pytorch和Torchvision

久未

已于 2023-08-17 14:37:40 修改

阅读量1.2k

点赞数 8

分类专栏： AI 文章标签： pytorch 人工智能 python

于 2023-08-17 14:34:15 首次发布

本文链接：https://blog.csdn.net/weixin_44604409/article/details/132334866

版权

AI 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

NVIDIA ORIN配置Pytorch和Torchvision

1、环境介绍
2、安装Pytorch
3、安装torchvision

1、环境介绍

NVIDIA AGX ORIN 64GB版本，通过SDK配置了JetPack 5.1.2(CUDA 11.4)，希望安装pytorch和torchvision。头一次配置移动开发平台，作为记录与参考，欢迎讨论与交流。

2、安装Pytorch

2.1 确定Pytorch版本

虽然https://elinux.org/Jetson_Zoo#PyTorch_.28Caffe2.29显示:JetPack 5.1可以使用v1.14.0、v2.0.0和v2.1.0中的任意一个，但是我尝试过Pytorch-v2.0.0并按照章节3的步骤安装torchvision，实际运行代码显示torchvision没找到，有知道原因的友友欢迎讨论交流。
根据https://developer.download.nvidia.com/compute/redist/jp/v$JP_VERSION/pytorch/$PYT_VERSION 页面指导，访问https://developer.download.nvidia.cn/compute/redist/jp/v512/，JetPack 5.1.2推荐的pytorch版本为v2.1.0，下载其whl文件（网址中v512表示JetPack 5.1.2）

2.2 安装Pytorch

使用pip install -tab(本地whl文件地址)安装Pytorch，建议使用3.8版本的python，版本过高可能会出错。

2.3 验证Pytorch

import torch
print(torch.cuda.is_available())
# 为 True表示安装成功

3、安装torchvision

3.0 踩坑一

使用https://pypi.org/project/torchvision/0.15.1/中的命令pip install torchvision==0.15.X安装会检测到torch版本不符合，会~~贴心地~~ 帮你删掉你刚安装好的Pytorch，再安装标准版（CPU）的pytorch。
仅做示意图，torch版本参考正文

3.0 踩坑二

安装torchvision时，参考了一些不同来源的资料，其中使用了sudo python3 setup.py install --user命令，将torchvision安装到root环境下，在普通环境下pip list 不显示torchvision，python代码也找不到torchvision库，运行时会出错。（https://zhuanlan.zhihu.com/p/425156043，注：某些文章中说将anaconda安装在root下，也许跟此处将torchvision安装在root下有相关性？Anyway，我的Anaconda安装在普通环境，因此需要按照以下步骤才能正确配置好环境。）

3.1 确定torchvision版本

根据NVIDIA论坛：https://forums.developer.nvidia.com/t/pytorch-for-jetson/72048页面描述：
torchvision版本及安装指令
猜测Pytorch v2.1.0对应的torchvision版本为v0.15.2（此处版本选择为实践结果，并非本人实际理论水平，即：JetPack5.1.2、CUDA 11.4环境下，使用Pytorch-v2.1.0和torchvision-v0.15.2能够运行）。

3.2 安装torchvision

上述英伟达论坛网址提供以下命令，仅供参考，对我个人有效：

sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev 
sudo apt-get install libavformat-dev libswscale-dev
git clone --branch  v0.15.2 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.15.2
python3 setup.py install --user

如果出现以下错误：

/usr/local/cuda/bin/nvcc: No such file or directory

根据更多参考链接中解决方案：

将~/.bashrc中的CUDA_HOME改为nvcc的路径：export CUDA_HOME=/usr/local/cuda
然后source ~/.bashrc

更多参考链接

3.3 验证torchvision

import torchvision
print(torchvision.__version__)
-----------------------------------------------------------
输出结果为：
(base) ****@ubuntu:~/****/****/****$ python
Python 3.8.17 (default, Jul  5 2023, 20:40:03) 
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torchvision
print(torchvision.__version__)
>>> print(torchvision.__version__)
0.15.2